在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实例：
--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
本文为原创文章，版权归知行编程网所有，欢迎分享本文，转载请保留出处！
内容反馈