こんばんは、Dajiroです。今回はGoogle Colabratory(以下、Colab)におけるPyTorchの使い方についてご紹介します。ColabといえばGoogle社が無料で提供しているノートブック形式のPython計算環境です。通常のCPUに加え、GPUとTPUといった機械学習向けの計算環境も使えるため、手っ取り早く爆速で計算を回すのには最強の環境と言えると思います。使い方は簡単、以下のリンクに飛ぶだけです。恐ろしい・・・
colab.research.google.com
さて、Colabの最大の特徴の１つは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を選択します。
セルに以下を入力し、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の挙動が理解できていません。
おわりに
最近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
