本文經(jīng)AI新媒體量子位（公眾號ID:QbitAI）授權(quán)轉(zhuǎn)載，轉(zhuǎn)載請聯(lián)系出處。
Facebook在PyTorch開發(fā)者大會(huì)上正式推出了PyTorch 1.3，并宣布了對谷歌云TPU的全面支持，而且還可以在Colab中調(diào)用云TPU。
之前機(jī)器學(xué)習(xí)開發(fā)者雖然也能在Colab中使用PyTorch，但是支持云TPU還是第一次，這也意味著你不需要購買昂貴的GPU，可以在云端訓(xùn)練自己的模型。
而且如果你是谷歌云平臺(tái)（Google Cloud Platform）的新注冊用戶，還能獲得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ū)域。
如果要對Resnet50在真實(shí)數(shù)據(jù)上進(jìn)行訓(xùn)練，需要選擇具有最多CPU數(shù)量的機(jī)器類型。為了獲得最佳效果，請選擇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
官方建議初次運(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)練：
用Docker訓(xùn)練：
在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)練完成后請記得刪除虛擬機(jī)和TPU。
性能比GPU提升4倍
訓(xùn)練完成后，我們就可以在Colab中導(dǎo)入自己的模型了。
打開notebook文件，在菜單欄的Runtime中選擇Change runtime type，將硬件加速器的類型改成TPU。
先運(yùn)行下面的代碼單元格，確?？梢栽L問Colab上的TPU：
然后在Colab中安裝兼容PyTorch/TPU組件：
接下來就可以導(dǎo)入你要訓(xùn)練好的模型和需要進(jìn)行推理的圖片了。
在PyTorch上使用TPU對性能的提升到底有多明顯呢？官方選用了v2-8的一個(gè)核心，即1/8 TPU的情形，與使用英偉達(dá)Tesla K80 GPU進(jìn)行對比，實(shí)測顯示推理時(shí)間大大縮短，性能約有4倍左右的提升。
GitHub地址：
https://github.com/pytorch/xla/tree/master/contrib/colab
【責(zé)任編輯：張燕妮 TEL：（010）68476606】