设为首页 - 加入收藏 双鸭山365棋牌+如何兑换_m.365you.com指尖棋牌_365棋牌什么好玩 (http://www.0469zz.com)- 国内知名站长资讯网站,提供最新最全的站长资讯,创业经验,网站建设等!
热搜: 百度 创业 网站 小米
当前位置: 首页 > 运营中心 > 建站资源 > 经验 > 正文

PyTorch终于能用上谷歌云TPU,推理性能提升4倍,该如何薅羊毛?

发布时间:2019-10-15 20:44 所属栏目:[经验] 来源:晓查
导读:本文经AI新媒体量子位(公众号ID:QbitAI)授权转载,转载请联系出处。 Facebook在PyTorch开发者大会上正式推出了PyTorch 1.3,并宣布了对谷歌云TPU的全面支持,而且还可以在Colab中调用云TPU。 之前机器学习开发者虽然也能在Colab中使用PyTorch,但是支持

本文经AI新媒体量子位(公众号ID:QbitAI)授权转载,转载请联系出处。

Facebook在PyTorch开发者大会上正式推出了PyTorch 1.3,并宣布了对谷歌云TPU的全面支持,而且还可以在Colab中调用云TPU。

之前机器学习开发者虽然也能在Colab中使用PyTorch,但是支持云TPU还是第一次,这也意味着你不需要购买昂贵的GPU,可以在云端训练自己的模型。

而且如果你是谷歌云平台(Google Cloud Platform)的新注册用户,还能获得300美元的免费额度。

PyTorch终于能用上谷歌云TPU,推理性能提升4倍,该如何薅羊毛?

现在PyTorch官方已经在Github上给出示例代码,教你如何免费使用谷歌云TPU训练模型,然后在Colab中进行推理。

训练ResNet-50

PyTorch先介绍了在云TPU设备上训练ResNet-50模型的案例。如果你要用云TPU训练其他的图像分类模型,操作方式也是类似的。

在训练之前,我们先要转到控制台创建一个新的虚拟机实例,指定虚拟机的名称和区域。

PyTorch终于能用上谷歌云TPU,推理性能提升4倍,该如何薅羊毛?

如果要对Resnet50在真实数据上进行训练,需要选择具有最多CPU数量的机器类型。为了获得最佳效果,请选择n1-highmem-96机器类型。

然后选择Debian GNU/Linux 9 Stretch + PyTorch/XLA启动盘。如果打算用ImageNet真实数据训练,需要至少300GB的磁盘大小。如果使用假数据训练,默认磁盘大小只要20GB。

创建TPU

  1. 转到控制台中创建TPU。
  2. 在“Name”中指定TPU Pod的名称。
  3. 在“Zone”中指定云TPU的区域,确保它与之前创建的虚拟机在同一区域中。
  4. 在“ TPU Type”下,选择TPU类型,为了获得最佳效果,请选择v3-8TPU(8个v3)。
  5. 在“ TPU software version”下,选择最新的稳定版本。
  6. 使用默认网络。
  7. 设置IP地址范围,例如10.240.0.0。

官方建议初次运行时使用假数据进行训练,因为fake_data会自动安装在虚拟机中,并且只需更少的时间和资源。你可以使用conda或Docker进行训练。

在fake_data上测试成功后,可以开始尝试用在ImageNet的这样实际数据上进行训练。

用conda训练:

  1. #?Fill?in?your?the?name?of?your?VM?and?the?zone.?
  2. $?gcloud?beta?compute?ssh?"your-VM-name"?--zone?"your-zone".?
  3. (vm)$?export?TPU_IP_ADDRESS=your-ip-address?
  4. (vm)$?export?XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"?
  5. (vm)$?ulimit?-n?10240?
  6. (vm)$?conda?activate?torch-xla-0.5?
  7. (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训练:

  1. #?Fill?in?your?the?name?of?your?VM?and?the?zone.?
  2. $?gcloud?beta?compute?ssh?"your-VM-name"?--zone?"your-zone".?
  3. (vm)$?export?TPU_IP_ADDRESS=your-ip-address?
  4. (vm)$?docker?run?--shm-size?128G?-v?~/imagenet:/tmp/imagenet?-e?XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"?gcr.io/tpu-pytorch/xla:r0.5?python3?pytorch/xla/test/test_train_imagenet.py?--model=resnet50?--num_epochs=90?--num_workers=64?--log_steps=200?--datadir=/tmp/imagenet?

在n1-highmem-96的虚拟机上选用完整v3-8 TPU进行训练,第一个epoch通常需要约20分钟,而随后的epoch通常需要约11分钟。该模型在90个epoch后达到约76%的top-1准确率。

为了避免谷歌云后续进行计费,在训练完成后请记得删除虚拟机和TPU。

性能比GPU提升4倍

训练完成后,我们就可以在Colab中导入自己的模型了。

打开notebook文件,在菜单栏的Runtime中选择Change runtime type,将硬件加速器的类型改成TPU。

PyTorch终于能用上谷歌云TPU,推理性能提升4倍,该如何薅羊毛?

先运行下面的代码单元格,确保可以访问Colab上的TPU:

  1. import?os?
  2. assert?os.environ[‘COLAB_TPU_ADDR’],?‘Make?sure?to?select?TPU?from?Edit?>?Notebook?settings?>?Hardware?accelerator’?

然后在Colab中安装兼容PyTorch/TPU组件:

  1. DIST_BUCKET="gs://tpu-pytorch/wheels"?
  2. TORCH_WHEEL="torch-1.15-cp36-cp36m-linux_x86_64.whl"?
  3. TORCH_XLA_WHEEL="torch_xla-1.15-cp36-cp36m-linux_x86_64.whl"?
  4. TORCHVISION_WHEEL="torchvision-0.3.0-cp36-cp36m-linux_x86_64.whl"?
  5. #?Install?Colab?TPU?compat?PyTorch/TPU?wheels?and?dependencies?
  6. !pip?uninstall?-y?torch?torchvision?
  7. !gsutil?cp?"$DIST_BUCKET/$TORCH_WHEEL"?.?
  8. !gsutil?cp?"$DIST_BUCKET/$TORCH_XLA_WHEEL"?.?
  9. !gsutil?cp?"$DIST_BUCKET/$TORCHVISION_WHEEL"?.?
  10. !pip?install?"$TORCH_WHEEL"?
  11. !pip?install?"$TORCH_XLA_WHEEL"?
  12. !pip?install?"$TORCHVISION_WHEEL"?
  13. !sudo?apt-get?install?libomp5?

接下来就可以导入你要训练好的模型和需要进行推理的图片了。

【免责声明】本站内容转载自互联网,其相关言论仅代表作者个人观点绝非权威,不代表本站立场。如您发现内容存在版权问题,请提交相关链接至邮箱:bqsm@foxmail.com,我们将及时予以处理。

网友评论
推荐文章