ころがる狸

ころがる狸のデータ解析ブログ

【PyTorch×TPU】Google ColabでPyTorchを使ってみた

こんばんは、Dajiroです。今回はGoogle Colabratory(以下、Colab)におけるPyTorchの使い方についてご紹介します。ColabといえばGoogle社が無料で提供しているノートブック形式のPython計算環境です。通常のCPUに加え、GPUとTPUといった機械学習向けの計算環境も使えるため、手っ取り早く爆速で計算を回すのには最強の環境と言えると思います。使い方は簡単、以下のリンクに飛ぶだけです。恐ろしい・・・
colab.research.google.com

さて、Colabの最大の特徴の1つはTPU (Tensor Processing Unit)が使えることです。TPUはGoogle社が開発した機械学習に特化したプロセッサです。機械学習では32, 64ビット等の計算精度が求められないため、TPUは8/16ビットの演算器から構成されています。このように使用用途をあえて機械学習向けに絞ることによって、省エネと高い計算量の両立が可能となっています。Googleが開発したこともあり、TensorFlowとTPUの親和性は高いですがPyTorchでのTPUを使うことも可能です。以下ではColab上でのその使い方を見ていきます。

【目次】

PyTorchにおけるTPUの使い方

Colab上でPyTorchの学習済みモデルを用いた画像分類の転移学習を行いました。学習対象はImageNetV2という10000枚の画像からなる画像データセットを用いました。バッチサイズ32, 学習率0.00001(Adam), 学習済みモデル(wide resnet50_2)のパラメータは全て学習対象としています。コード全文はGitHubのレポジトリに保存しているため、以下では特に重要な部分のみ記載いたします。
github.com

PyTorch/XLAのインストール

線形代数計算の最適化コンパイラにXLAというものがあります。これはTensorFlowの計算最適化のために開発されたようですが、これをPyTorchにも使えるようにするPyTorch/XLAをインストールすることでTPUが使用可能になります。

まずは、Colab画面を開き編集⇒ノートブックの設定でTPUを選択します。

f:id:Dajiro:20200725104537p:plain
ColabでTPUを選択。

セルに以下を入力し、PyTorch/XLAをインストールします。

VERSION = "20200325"  #@param ["1.5" , "20200325", "nightly"]
!curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py
!python pytorch-xla-env-setup.py --version $VERSION

TPUの設定

ライブラリをimportして、TPUのデバイスを定義します。deviceにxla:1と設定されていれば準備完了です。

import torch_xla
import torch_xla.core.xla_model as xm
device = xm.xla_device()
print(device)        #xla:1

あとは一般的なPyTorchの書き方と同様に、.to(device)によってモデルや入力をTPUに載せます。注意点としては学習の更新を行う際に、xlmからオプティマイザを呼び出す必要があります。ざっくりとした学習用コードはこんな感じですが、パラメータ更新部以外は普通のPyTorchにおける学習法とまったく同じです。

#学習済みモデル定義
class WResNet(nn.Module):
    def __init__(self, n_out):
        super(WResNet, self).__init__()
        #wide resnet50_2を使用。パラメータ数は約7千万。
        self.resnet = models.wide_resnet50_2(pretrained=True)
        self.resnet.fc = nn.Linear(2048, n_out)

    def forward(self, x):
        return self.resnet(x)
# 損失関数
criterion = nn.CrossEntropyLoss()
#オプティマイザ
optimizer = optim.Adam(model.parameters(), lr=lr, betas=(beta1, 0.999))
#モデル定義しTPUデバイスに載せる
model = WResNet(n_class)
model.to(device)
#学習実行部分
for epoch in range(epochs):
    for inputs, labels in dataloaders
         #入出力をTPUデバイスに載せる
         labels = labels.to(device)
         inputs = inputs.to(device)
         outputs = model(inputs)       
         loss = criterion(outputs, labels)
         _, preds = torch.max(outputs, 1)
         loss.backward()
         #xlmからパラメータの更新を行う
         xm.optimizer_step(optimizer, barrier=True)

GPUとの比較

上記学習のTPUとGPUにおける計算速度の比較を行いました。学習に時間がかかるため5エポックで計算を打ち切りましたが、やや意外な結果に。計算条件は同じでそれぞれ複数回計算を回しましたが、TPUでは7分58秒、12分35秒、8分35秒と計算速度にばらつきが出ました。一方GPUでは11分15秒、11分8秒と比較的安定した結果に。最速はTPUの1回目の試行でしたが、どうして同じTPU計算でこれほど差が出るのでしょう?イマイチColabにおける計算資源の割り振りとか、TPUの挙動が理解できていません。

f:id:Dajiro:20200725191235p:plain
TPUとGPUの速度比較。それぞれ複数回計算(学習データはシャッフル)。

おわりに

最近PyTorchの使用頻度が増えたためColabでのTPU計算を行ってみましたが、非常に使い勝手が良いと感じました。速度はGPUよりもTPUの方が早そうですが、結果がやや安定しません。計算条件や使うモデルの構造によっても結果は変わってくるのかもしれません。GPUとTPUの速度比較は以下のQiita記事が詳しいためご覧ください。
Google ColabのTPUで対GPUの最速に挑戦する - Qiita

また、今回はPyTorchが公式で提供している学習済みモデルを使いましたが、TIMMなどの外部ライブラリのモデルを使うとTPUのエラーが出るという事象に遭遇しています。これはTorchScriptのエラーのようですが、結局対処できませんでした。もっとPyTorch×TPUの知見が欲しいところです。今回の記事がどなたかの参考になれば幸いです。

/usr/local/lib/python3.6/dist-packages/timm/models/layers/activations_me.py in forward(ctx, x)
     35     def forward(ctx, x):
     36         ctx.save_for_backward(x)
---> 37         return swish_jit_fwd(x)
     38 
     39     @staticmethod

RuntimeError: Unknown device