本文经AI新媒体量子比特（公众号ID：QbitAI）授权转载，转载请联系出处。Facebook在PyTorch开发者大会上正式推出了PyTorch1.3，并宣布全面支持GoogleCloudTPU，也可以在Colab中调用CloudTPU。虽然之前机器学习开发者也可以在Colab中使用PyTorch，但支持云TPU还是第一次。这也意味着你不需要购买昂贵的GPU，你可以在云端训练自己的模型。而如果你是谷歌云平台（GoogleCloudPlatform）的新用户，你还可以获得300美元的免费赠金。现在PyTorch官方在Github上给出了示例代码，教大家免费使用GoogleCloudTPU训练模型，然后在Colab中进行推理。训练ResNet-50PyTorch首先介绍在云端TPU设备上训练ResNet-50模型的案例。如果想使用CloudTPU训练其他图像分类模型，操作方法类似。在训练之前，我们要到控制台创建一个新的虚拟机实例，指定虚拟机的名称和区域。如果要在真实数据上训练Resnet50，需要选择CPU数量最多的机器类型。为获得最佳结果，请选择n1-highmem-96机器类型。然后选择DebianGNU/Linux9Stretch+PyTorch/XLA启动盘。如果您计划使用ImageNet真实数据进行训练，则需要至少300GB的磁盘大小。如果用假数据训练，默认磁盘大小只有20GB。创建TPU进入控制台创建TPU。在“名称”中指定TPUPod的名称。在“Zone”中指定云TPU的区域，确保它与之前创建的虚拟机位于同一区域。在TPU类型下，选择TPU类型，为获得最佳结果，请选择v3-8TPU(8v3s)。在“TPU软件版本”下，选择最新的稳定版本。使用默认网络。设置IP地址范围，例如10.240.0.0。官方推荐第一次运行时使用fakedata进行训练，因为fake_data是自动安装在虚拟机中的，需要的时间和资源更少。您可以使用conda或Docker进行训练。在fake_data上测试成功后，就可以开始尝试在ImageNet的实际数据上训练了。使用conda进行训练：#FillinyourthenameofyourVMandthezone.$gcloudbetacomputessh"your-VM-name"--zone"your-zone".(vm)$exportTPU_IP_ADDRESS=your-ip-address(vm)$exportXRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"(vm)$ulimit-n10240(vm)$condaactivatetorch-xla-0.5(torch-xla-0.5)$python/usr/share/torch-xla-0.5/pytorch/xla/test/test_train_imagenet.py--datadir=~/imagenet--model=resnet50--num_epochs=90--num_workers=64--batch_size=128--log_steps=200使用Docker进行训练：#FillinyourthenameofyourVMandthezone.$gcloudbetacomputessh"your-VM-name"--zone"your-zone".(vm)$exportTPU_IP_ADDRESS=你的IP地址(vm)$dockerrun--shm-size128G-v~/imagenet:/tmp/imagenet-eXRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"gcr。io/tpu-pytorch/xla:r0.5python3pytorch/xla/test/test_train_imagenet.py--model=resnet50--num_epochs=90--num_workers=64--log_steps=200--datadir=/tmp/imagenet在n1-选择完整的v3-8TPU在highmem-96虚拟机上进行训练。第一个epoch通常需要20分钟左右，随后的epoch通常需要大约11分钟，模型在90个epoch后达到大约76%的top-1准确率。为避免后续被谷歌云计费，请记得在训练完成后删除虚拟机和TPU。性能是GPU的4倍。训练完成后，我们就可以在Colab中导入自己的模型了。打开notebook文件，在菜单栏的Runtime中选择Changeruntimetype，将硬件加速器类型改为TPU。首先运行以下代码单元以确保您可以在Colab上访问TPU：importosassertos.environ['COLAB_TPU_ADDR'],'MakesuretoselectTPUfromEdit>Notebooksettings>Hardwareaccelerator'然后在Colab中安装兼容的PyTorch/TPU组件：DIST_BUCKET="gs://tpu-pytorch/wheels"TORCH_WHEEL="torch-1.15-cp36-cp36m-linux_x86_64.whl"TORCH_XLA_WHEEL="torch_xla-1.15-cp36-cp36m-linux_x86_64.whl"TORCHVISION_WHEEL="torchvision-0.6.0-cp36-4.x386m.x386"#InstallColabTPUcompatPyTorch/TPUwheelsanddependencies!pipuninstall-ytorchtorchvision!gsutilcp"$DIST_BUCKET/$TORCH_WHEEL".!gsutilcp"$DIST_BUCKET/$TORCH_XLA_WHEEL".!gsutilcp"$DIST_BUCKET/$TORCHVISION_WHEEL".!pipinstall"$TORCH_WHEEL"$HE_WHEEL"CHTOR!pip"!pipinstall"$TORCHVISION_WHEEL"!sudoapt-getinstalllibomp5接下来就可以导入你要训练的模型和需要推断的图片了。在PyTorch上使用TPU的性能提升有多大？官方选择了一个v2-8核心，即1/8TPU的情况下，与使用NVIDIATeslaK80GPU相比，实测显示推理时间大大缩短，性能提升约4次。GitHub地址：https://github.com/pytorch/xla/tree/master/contrib/colab
