在2019年PyTorch開發者大會上，Facebook，Google和Salesforce Research聯合宣布啟動PyTorch-TPU項目。項目的目標是在保持PyTorch的靈活性的同時讓社區盡可能容易地利用云TPU提供的高性能計算。團隊創建了PyTorch/XLA這個repo，它可以讓使PyTorch連接到云TPU并使用TPU內核，同時Colab也支持在云TPU上使用PyTorch/XLA。現在，PyTorch/XLA 已在Google Cloud上達到通用標準(GA)，并為我們提供了許多方便易用的接口。趕緊來看看PyTorch/XLA都有哪些功能吧

GA版本有哪些新功能？

借助PyTorch/XLA GA，云TPU正式支持PyTorch 1.6。其他值得注意的新功能包括：

• 支持層內模型并行性：現在可以在reduce的時候在多組tensor上使用多種運算。添加了更多的通信原語從而可以實現有趣的應用，例如把嵌入詞向量詞分布到多個TPU內核上；
• 額外的XLA運算：隨著PyTorch/XLA在越來越廣泛的新模型中被使用，用戶要求將PyTorch運算映射到XLA。對此，從beta(1.5)版本開始，我們已經為Replication_pad1d，replication_pad2d，max_unpool2d，max_unpool3d等運算加入了低配XLA；
• 在Colab/Kaggle上更好的體驗：現在，不再需要在Colab/Kaggle上運行env-setup.py腳本，就可以開始訓練倆。
• 對深度學習VM映像的支持：Google Cloud 平臺提供了一組深度學習虛擬機(DLVM)映像，配置好了各種常用深度學習框架(包括PyTorch)和所需的一切。PyTorch/XLA 1.6現在已預安裝在DLVM中，并針對云TPU進行了優化。官方的PyTorch 1.6也預裝在相同的Conda環境中。

詳細使用方法參照用戶指南^[1]。

GA版本支持哪些模型？

PyTorch/XLA已被用于在云TPU上訓練眾多深度學習模型。包括：

• 圖像分類任務(ImageNet)：Torchvision的ResNet-50；
• 翻譯任務(WMT-18 en-de)：Fairseq Transformer；
• 通用的語言模型的預訓練和微調(GLUE等)：HuggingFace(BERT，DistilBERT，RoBERTa，XLNet等)和Fairseq RoBERTa；
• 深度學習推薦模型(1.6版的新增功能)：DLRM。

在大多數情況下，在云TPU上訓練這些模型幾乎不需要更改代碼。可以參考官方教程^[2]獲得上述模型的詳細教程，也可以參考PyTorch / XLA GitHub^[3]代碼庫找到訓練的其他模型架構的示例。

PyTorch/XLA是如何工作的？

PyTorch/XLA使用“lazy tensor”進行抽象。使用lazy tensor時，運算的evaluation會被推遲到這個運算被訪問之前。此時，運算被描述為中間表示圖，當我們需要運算結果時(訪問)，就通過XLA編譯這些中間表示圖，并發送到TPU內核以執行。這種XLA編譯對CPU和GPU都支持。其他技術詳細信息，可以參考GitHub^[4]。

升級到云TPU要改哪些代碼？

首先需要創建一個配有PyTorch / XLA映像的Google Compute Engine虛擬機和一個云TPU實例。創建好虛擬機和云TPU實例后，需要配置conda環境并設置XRT_TPU_CONFIG環境變量指向云TPU實例：

export?XRT_TPU_CONFIG="tpu_worker;0;:8470"

現在就可以開始在云TPU上訓練模型了！

實際代碼中需要改動的地方包括：

• 調取相關代碼包import torch_xla
• 設置訪問XLA設備抽象的方法以及并行數據加載器。
• 使用xm.optimizer_step(optimizer)在后臺進行reduce。

可以參考英文博客原文^[5]中的代碼實例看看具體的區別。

在Cloud TPU Pod上訓練

PyTorch/XLA支持將剛剛在單個云TPU上執行的訓練擴展到整個云TPU Pod或者任意Pod切片,只需要使用xla_dist包裝就可以了：

?python?-m?torch_xla.distributed.xla_dist?\
??????--tpu=$TPU_NAME?\
??????--conda-env=torch-xla-1.6?\
??????--env?ANY_ENV_VAR=VALUE?\
??????--?\
??????python?/path/to/your/code.py?--train_arg1?\
????????--train_arg2?...

官方提供的Colab notebook^[6]寫得非常詳細，現在就開始探索吧！

更多詳情請參考英文博客原文^[7]以及Google官方博文^[8]

參考文獻

[1]用戶指南: https://docs.google.com/document/d/1RzNgpTK4wESImhIwgMbknqn9xVNUFtCTCknNfONlMlo/edit

[2]官方教程: https://cloud.google.com/tpu/docs/tutorials/

[3]PyTorch / XLA GitHub: https://github.com/pytorch/xla

[4]GitHub: https://github.com/pytorch/xla/blob/master/API_GUIDE.md#xla-tensor-deep-dive

[5]英文博客原文: https://medium.com/pytorch/pytorch-xla-is-now-generally-available-on-google-cloud-tpus-f9267f437832

[6]Colab notebook: https://github.com/pytorch/xla/tree/master/contrib/colab

[7]英文博客原文: https://medium.com/pytorch/pytorch-xla-is-now-generally-available-on-google-cloud-tpus-f9267f437832

[8]Google官方博文: https://cloud.google.com/blog/products/ai-machine-learning/pytorch-is-now-ga-on-google-cloud-tpu

---

推薦閱讀

(點擊標題可跳轉閱讀)

干貨 | 公眾號歷史文章精選

我的深度學習入門路線

我的機器學習入門路線圖

重磅！

AI有道年度技術文章電子版PDF來啦！

掃描下方二維碼，添加?AI有道小助手微信，可申請入群，并獲得2020完整技術文章合集PDF(一定要備注：入群?+ 地點 + 學校/公司。例如：入群+上海+復旦。?

長按掃碼，申請入群

(添加人數較多，請耐心等待)

感謝你的分享，點贊，在看三連↓

總結

以上是生活随笔為你收集整理的conda如何升级pytorch_Google Cloud TPUs 支持 Pytorch 框架啦！的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。