曉查 發自 凹非寺量子位 報導 | 公眾號 QbitAI
昨天，Facebook在PyTorch開發者大會上正式推出了PyTorch 1.3，並宣布了對谷歌雲TPU的全面支持，而且還可以在Colab中調用雲TPU。
之前機器學習開發者雖然也能在Colab中使用PyTorch，但是支持雲TPU還是第一次，這也意味著你不需要購買昂貴的GPU，可以在雲端訓練自己的模型。
而且如果你是谷歌雲平臺（Google Cloud Platform）的新註冊用戶，還能獲得300美元的免費額度。
現在PyTorch官方已經在Github上給出示例代碼，教你如何免費使用谷歌雲TPU訓練模型，然後在Colab中進行推理。
訓練ResNet-50
PyTorch先介紹了在雲TPU設備上訓練ResNet-50模型的案例。如果你要用雲TPU訓練其他的圖像分類模型，操作方式也是類似的。
在訓練之前，我們先要轉到控制臺創建一個新的虛擬機實例，指定虛擬機的名稱和區域。
如果要對Resnet50在真實數據上進行訓練，需要選擇具有最多CPU數量的機器類型。為了獲得最佳效果，請選擇n1-highmem-96機器類型。
然後選擇Debian GNU/Linux 9 Stretch + PyTorch/XLA啟動盤。如果打算用ImageNet真實數據訓練，需要至少300GB的磁碟大小。如果使用假數據訓練，默認磁碟大小隻要20GB。
創建TPU
轉到控制臺中創建TPU。在「Name」中指定TPU Pod的名稱。在「Zone」中指定雲TPU的區域，確保它與之前創建的虛擬機在同一區域中。在「 TPU Type」下，選擇TPU類型，為了獲得最佳效果，請選擇v3-8TPU（8個v3）。在「 TPU software version」下，選擇最新的穩定版本。使用默認網絡。設置IP位址範圍，例如10.240.0.0。官方建議初次運行時使用假數據進行訓練，因為fake_data會自動安裝在虛擬機中，並且只需更少的時間和資源。你可以使用conda或Docker進行訓練。
在fake_data上測試成功後，可以開始嘗試用在ImageNet的這樣實際數據上進行訓練。
用conda訓練：
# 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訓練：
# 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的虛擬機上選用完整v3-8 TPU進行訓練，第一個epoch通常需要約20分鐘，而隨後的epoch通常需要約11分鐘。該模型在90個epoch後達到約76％的top-1準確率。
為了避免谷歌雲後續進行計費，在訓練完成後請記得刪除虛擬機和TPU。
性能比GPU提升4倍
訓練完成後，我們就可以在Colab中導入自己的模型了。
打開notebook文件，在菜單欄的Runtime中選擇Change runtime type，將硬體加速器的類型改成TPU。
先運行下面的代碼單元格，確保可以訪問Colab上的TPU：
import osassert 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
接下來就可以導入你要訓練好的模型和需要進行推理的圖片了。
在PyTorch上使用TPU對性能的提升到底有多明顯呢？官方選用了v2-8的一個核心，即1/8 TPU的情形，與使用英偉達Tesla K80 GPU進行對比，實測顯示推理時間大大縮短，性能約有4倍左右的提升。
GitHub地址：https://github.com/pytorch/xla/tree/master/contrib/colab