炼数成金 门户 大数据 云计算 查看内容

Google Cloud TPUs支持Pytorch框架啦!

2020-11-3 16:08| 发布者: 炼数成金_小数| 查看: 67790| 评论: 0|原作者: Sherry|来自: 夕小瑶的卖萌屋

摘要: 在2019年PyTorch开发者大会上,Facebook,Google和Salesforce Research联合宣布启动PyTorch-TPU项目。项目的目标是在保持PyTorch的灵活性的同时让社区尽可能容易地利用云TPU提供的高性能计算。团队创建了PyTorch/XLA ...
在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;<TPU_IP_ADDRESS>: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

声明:文章收集于网络,版权归原作者所有,为传播信息而发,如有侵权,请联系小编删除,谢谢!

欢迎加入本站公开兴趣群
软件开发技术群
兴趣范围包括:Java,C/C++,Python,PHP,Ruby,shell等各种语言开发经验交流,各种框架使用,外包项目机会,学习、培训、跳槽等交流
QQ群:26931708

Hadoop源代码研究群
兴趣范围包括:Hadoop源代码解读,改进,优化,分布式系统场景定制,与Hadoop有关的各种开源项目,总之就是玩转Hadoop
QQ群:288410967 

鲜花

握手

雷人

路过

鸡蛋

相关阅读

最新评论

热门频道

  • 大数据
  • 商业智能
  • 量化投资
  • 科学探索
  • 创业

即将开课

热门文章

     

    GMT+8, 2021-8-1 11:15 , Processed in 0.171657 second(s), 24 queries .