曉查 發(fā)自 凹非寺
量子位 報(bào)道 | 公眾號(hào) QbitAI

Facebook在PyTorch開發(fā)者大會(huì)上正式推出了PyTorch 1.3，并宣布了對(duì)谷歌云TPU的全面支持，而且還可以在Colab中調(diào)用云TPU。

之前機(jī)器學(xué)習(xí)開發(fā)者雖然也能在Colab中使用PyTorch，但是支持云TPU還是第一次，這也意味著你不需要購買昂貴的GPU，可以在云端訓(xùn)練自己的模型。

而且如果你是谷歌云平臺(tái)(Google Cloud Platform)的新注冊(cè)用戶，還能獲得300美元的免費(fèi)額度。

現(xiàn)在PyTorch官方已經(jīng)在Github上給出示例代碼，教你如何免費(fèi)使用谷歌云TPU訓(xùn)練模型，然后在Colab中進(jìn)行推理。

訓(xùn)練ResNet-50

PyTorch先介紹了在云TPU設(shè)備上訓(xùn)練ResNet-50模型的案例。如果你要用云TPU訓(xùn)練其他的圖像分類模型，操作方式也是類似的。

在訓(xùn)練之前，我們先要轉(zhuǎn)到控制臺(tái)創(chuàng)建一個(gè)新的虛擬機(jī)實(shí)例，指定虛擬機(jī)的名稱和區(qū)域。

如果要對(duì)Resnet50在真實(shí)數(shù)據(jù)上進(jìn)行訓(xùn)練，需要選擇具有最多CPU數(shù)量的機(jī)器類型。為了獲得最佳效果，請(qǐng)選擇n1-highmem-96機(jī)器類型。

然后選擇Debian GNU/Linux 9 Stretch + PyTorch/XLA啟動(dòng)盤。如果打算用ImageNet真實(shí)數(shù)據(jù)訓(xùn)練，需要至少300GB的磁盤大小。如果使用假數(shù)據(jù)訓(xùn)練，默認(rèn)磁盤大小只要20GB。

創(chuàng)建TPU

轉(zhuǎn)到控制臺(tái)中創(chuàng)建TPU。

在“Name”中指定TPU Pod的名稱。

在“Zone”中指定云TPU的區(qū)域，確保它與之前創(chuàng)建的虛擬機(jī)在同一區(qū)域中。

在“ TPU Type”下，選擇TPU類型，為了獲得最佳效果，請(qǐng)選擇v3-8TPU(8個(gè)v3)。

在“ TPU software version”下，選擇最新的穩(wěn)定版本。

使用默認(rèn)網(wǎng)絡(luò)。

設(shè)置IP地址范圍，例如10.240.0.0。

官方建議初次運(yùn)行時(shí)使用假數(shù)據(jù)進(jìn)行訓(xùn)練，因?yàn)閒ake_data會(huì)自動(dòng)安裝在虛擬機(jī)中，并且只需更少的時(shí)間和資源。你可以使用conda或Docker進(jìn)行訓(xùn)練。

在fake_data上測試成功后，可以開始嘗試用在ImageNet的這樣實(shí)際數(shù)據(jù)上進(jìn)行訓(xùn)練。

用conda訓(xùn)練：

#?Fill?in?your?the?name?of?your?VM?and?the?zone.
$?gcloud?beta?compute??ssh?"your-VM-name"?--zone?"your-zone".
(vm)$?export?TPU_IP_ADDRESS=your-ip-address
(vm)$?export?XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"
(vm)$?ulimit?-n?10240
(vm)$?conda?activate?torch-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訓(xùn)練：

#?Fill?in?your?the?name?of?your?VM?and?the?zone.
$?gcloud?beta?compute?ssh?"your-VM-name"?--zone?"your-zone".
(vm)$?export?TPU_IP_ADDRESS=your-ip-address
(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的虛擬機(jī)上選用完整v3-8 TPU進(jìn)行訓(xùn)練，第一個(gè)epoch通常需要約20分鐘，而隨后的epoch通常需要約11分鐘。該模型在90個(gè)epoch后達(dá)到約76％的top-1準(zhǔn)確率。

為了避免谷歌云后續(xù)進(jìn)行計(jì)費(fèi)，在訓(xùn)練完成后請(qǐng)記得刪除虛擬機(jī)和TPU。

性能比GPU提升4倍

訓(xùn)練完成后，我們就可以在Colab中導(dǎo)入自己的模型了。

打開notebook文件，在菜單欄的Runtime中選擇Change runtime type，將硬件加速器的類型改成TPU。

先運(yùn)行下面的代碼單元格，確保可以訪問Colab上的TPU：

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

然后在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.3.0-cp36-cp36m-linux_x86_64.whl"

#?Install?Colab?TPU?compat?PyTorch/TPU?wheels?and?dependencies
!pip?uninstall?-y?torch?torchvision
!gsutil?cp?"$DIST_BUCKET/$TORCH_WHEEL"?.
!gsutil?cp?"$DIST_BUCKET/$TORCH_XLA_WHEEL"?.
!gsutil?cp?"$DIST_BUCKET/$TORCHVISION_WHEEL"?.
!pip?install?"$TORCH_WHEEL"
!pip?install?"$TORCH_XLA_WHEEL"
!pip?install?"$TORCHVISION_WHEEL"
!sudo?apt-get?install?libomp5

接下來就可以導(dǎo)入你要訓(xùn)練好的模型和需要進(jìn)行推理的圖片了。

在PyTorch上使用TPU對(duì)性能的提升到底有多明顯呢？官方選用了v2-8的一個(gè)核心，即1/8 TPU的情形，與使用英偉達(dá)Tesla K80 GPU進(jìn)行對(duì)比，實(shí)測顯示推理時(shí)間大大縮短，性能約有4倍左右的提升。

GitHub地址：
https://github.com/pytorch/xla/tree/master/contrib/colab

作者系網(wǎng)易新聞·網(wǎng)易號(hào)“各有態(tài)度”簽約作者

— 完 —

大會(huì)啟幕！預(yù)見智能科技新未來

量子位MEET 2020智能未來大會(huì)啟幕，將攜手優(yōu)秀AI企業(yè)、杰出科研人員呈現(xiàn)一場高質(zhì)量行業(yè)盛會(huì)！詳情可點(diǎn)擊圖片：

榜單征集！三大獎(jiǎng)項(xiàng)，鎖定AI Top玩家

2019中國人工智能年度評(píng)選啟幕，將評(píng)選領(lǐng)航企業(yè)、商業(yè)突破人物、最具創(chuàng)新力產(chǎn)品3大獎(jiǎng)項(xiàng)，并于MEET 2020大會(huì)揭榜，歡迎優(yōu)秀的AI公司掃碼報(bào)名！

量子位?QbitAI · 頭條號(hào)簽約作者

?'?' ? 追蹤AI技術(shù)和產(chǎn)品新動(dòng)態(tài)

喜歡就點(diǎn)「好看」吧 !?

總結(jié)

以上是生活随笔為你收集整理的pytorch resnet50_PyTorch终于能用上谷歌云TPU，推理性能提升4倍，我们该如何薅羊毛？...的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。