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

作者: 晓查 2019-10-14 09:58:00

本文经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 

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

在PyTorch上使用TPU对性能的提升到底有多明显呢?官方选用了v2-8的一个核心,即1/8 TPU的情形,与使用英伟达Tesla K80 GPU进行对比,实测显示推理时间大大缩短,性能约有4倍左右的提升。

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

GitHub地址:

https://github.com/pytorch/xla/tree/master/contrib/colab

机器学习 人工智能 计算机
上一篇:2019年度十大Web开发趋势 下一篇:非监督学习最强攻略
评论
取消
暂无评论,快去成为第一个评论的人吧

更多资讯推荐

人工智能会给社会带来什么?一文看懂AI的变革 精选

人类科技迅猛发展的时代,人工智能以精准的算法和高效率的工作能力,在人类的生活与工作中起到愈加重要的作用,现在无论是手机、电脑等电子产品乃至于大型机器都涉及人工智能,为了让大家更好地了解人工智能,所以这次跟大家聊一聊人工智能的事儿~

科氪星球 ·  3天前
企业级AI如何开展?有Watson,无难题!

提到IBM的人工智能,当属Watson。2011年,Watson在美国智力问答电视节目《危险边缘》(Jeopardy)中亮相,并且一举打败了人类智力竞赛的冠军。从此,大家便知道了IBM有个叫Watson的人工智能平台。

鸢玮 ·  3天前
Reddit网友吵爆!算力和数据真能解决一切? 精选

近日,一位reddit用户发起一个讨论帖:如果我们只有更多的数据和计算能力而停止理论工作的发展,今天的哪些问题可以解决?哪些问题绝对无法解决?

张佳 ·  3天前
明年1月1日起,AI造假视频不得随意发布

新规来了。在这份刚刚印发的《网络音视频信息服务管理规定》中,四次直接提及深度学习,基本可以看做是针对AI造假视频的一次针对性管控。

问耕 ·  3天前
为物联网/人工智能项目寻找投资商的优秀方式

解决一个问题,并验证,而不是描述。许多物联网项目的演示方式,大大降低了投资者和客户的吸引力。这将导致项目永远无法启动,即使启动了,也很可能会失败。

Adam Dunkels ·  3天前
2020职场AI技能排行榜:TensorFlow热度飙升,Python最火,市场部也在学

近日,在线教育网站 Udemy 根据其学员的课程数据,制作了一份《2020 年职场学习趋势报告》,指出了哪些技能最受职场人关注。

佚名 ·  2019-11-29 13:57:55
无法击败 AI 李世石宣布退役

作为唯一一名曾经打败 AI 围棋程序 AlphaGo 的人类,韩国九段棋手李世石上周正式退出职业围棋。

佚名 ·  2019-11-29 08:00:36
百度与东软教育联合成立“东软百度人工智能学院” 破局AI“人才荒”

智能时代正在加速到来,而人工智能人才作为支撑AI发展的第一资源显得尤为重要。11月28日,百度与东软教育科技集团在京签署战略合作协议,并为合作设立的“东软百度人工智能学院”举行成立及授牌仪式。辽宁省教育厅党组成员、副厅长花蕾,国家开发银行大连分行行长张一帆,百度首席技术官王海峰,东软集团董事长兼CEO、东软教育科技集团董事长刘积仁等出席签约仪式。

佚名 ·  2019-11-28 18:19:06
Copyright©2005-2019 51CTO.COM 版权所有 未经许可 请勿转载