标签:git 目的 安装 replica 载器 code 环境 扩展 行数据
在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。其他值得注意的新功能包括:
GA版本支持哪些模型?
PyTorch/XLA已被用于在云TPU上训练众多深度学习模型。包括:
*图像分类任务(ImageNet)*:Torchvision的ResNet-50;
翻译任务(WMT-18 en-de)**:Fairseq Transformer;
*通用的语言模型的预训练和微调(GLUE等):HuggingFace(BERT,DistilBERT,RoBERTa,XLNet等)和Fairseq RoBERTa;
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上训练模型了!
实际代码中需要改动的地方包括:
在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]
Google Cloud TPUs支持Pytorch框架啦!
标签:git 目的 安装 replica 载器 code 环境 扩展 行数据
原文地址:https://blog.51cto.com/15061930/2567671