深層学習コンパイラTVMと主要深層学習フレームワークをColaboratoryで使い倒そう

2019/07/05 TVMのインストール方法を、ビルド済みバイナリを利用する方法に変更しました。
2019/04/04 タイトルを変更するとともに、内容を加筆修正しました

オンライン開発環境 Google Colaboratory

はじめに

はじめまして。皆さん、深層学習（Deep Learning）についてどんなイメージを持っていますか？「まさに技術的ブレイクスルーだ」というポジティブなイメージもあれば、「セットアップが大変で環境構築(GPU)にとにかくお金がかかる」というネガティブなイメージもあるでしょう。しかし、そんなネガティブなイメージを払拭するサービスがこのGoogle Colaboratoryです。

Google Colaboratory
https://colab.research.google.com/

これは簡単に説明すると、GCPのクラウドで執筆時（20180305）のレートで普通に借りたら、月額6万をオーバーしてしまうようなn1-highmem-2相当の環境を、連続起動時間が最長12時間という制約付きで、無料で借りられてしまうという衝撃的なサービスです。その上、面倒なCUDA、あるいはnvidia-dockerといったセットアップ作業も必要ありません。無料で、簡単に、今すぐDeep Learningを試せるのが、このColaboratoryなのです。

また2018/09/27頃からTPUが使えるようになりました。GPUとは浮動小数点演算の計算精度が変わるなど、注意点もありますが、Keras使用時でCNNが6倍高速になるケースがあるようです。詳細は、下記の記事をご覧ください。

Google ColabのTPUでResNetのベンチマークを取ってみた

事前準備

下記ベンチマークで検証した内容は、リンク先のColaboratory Noteから再現することができますが、１からノートを作る場合は以下の手順を踏んでください。

Colaboratoryを開く

新規ノートブックを作成

ノートブックを新規作成 > Python3の新しいノートブックを選ぶ。

GPUをアサイン

画面上部のメニュー ランタイム > ランタイムのタイプを変更で、ノートブックの設定を開く
ハードウェアアクセラレータにGPUを選択し、保存する

CUDAのバージョンをチェック

2019/04/04現在、CUDA 10.0がインストールされています。CUDAのバージョンによって使えるフレームワークのバージョンが変わってくるため、下記コマンドでCUDAのversionをチェックすることを強く推奨します。

!/usr/local/cuda/bin/nvcc --version
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2018 NVIDIA Corporation
Built on Sat_Aug_25_21:08:01_CDT_2018
Cuda compilation tools, release 10.0, V10.0.130

利用可能な深層学習フレームワーク

主要というのは主観になりますが、以下の深層学習フレームワークが利用できます。

Tensorflow / Keras

最初からインストールされています。2019/04/04時点ではtensorflow==1.13.1でした。頻繁にアップデートされるため、利用時のバージョンを記載した方がよいでしょう。もしくはpip install でバージョン指定して入れ直すという方法もありますが、CUDAのバージョンも変わるため永続的に動作する保証はありません。

PyTorch

最初からインストールされています。2019/04/04時点ではtorch==1.0.1.post2でした。

Chainer

最初からインストールされています。2019/04/04時点ではchainer==5.0.0でした。

MXNet

CUDAのバージョンにシビアなので、ご注意ください。

!pip install mxnet-cu100mkl

Caffe

取り上げた当時は無理と言われていたCaffeですが、なんと自前でビルドすることで実行できるようです。下記リンク先をご覧ください。
Build Caffe in Google Colaboratory: a free video card in the cloud

CNTK

!pip install cntk

Sonnet

最初からインストールされています。2019/04/04時点ではdm-sonnet==1.23でした。

Theano

最初からインストールされています。2019/04/04時点ではTheano==1.0.4でした。

フレームワークのベンチマーク

簡単なCNNによるディープラーニングライブラリ速度比較で紹介されている、CNNによる画像分類のベンチマークを実行してみました。

実行手順

!git clone https://github.com/daigo0927/compare-deeplibs.git
!rm compare-deeplibs/lap_record.csv
!cd compare-deeplibs;python train_tf.py -g 0 -e 5
!cd compare-deeplibs;python train_keras.py -g 0 -e 5
!cd compare-deeplibs;python train_sonnet.py -g 0 -e 5
!cd compare-deeplibs;python train_torch.py -g 0 -e 5
!cd compare-deeplibs;python train_chainer.py -g 0 -e 5

計算時間

	1エポック目[s]	2エポック目[s]
Tensorflow	15.03	14.33
Keras	17.08	15.66
Sonnet	16.53	15.77
PyTorch	3.24	2.54
Chainer	10.76	3.09

機械学習の可視化ツールを動かす

Tensorflow 2.0からはColaboratory上で直接TensorBoardが扱えます→公式サンプル

Tensorboardのような可視化ツールも、localtunnelを使うことで閲覧することができます。
先ほど紹介した、tensorflowによるベンチマークの実行結果を、tensorboardで表示してみましょう。以下、実行手順です。

! npm install -g localtunnel
LOG_DIR = './compare-deeplibs/model_tf/'
get_ipython().system_raw(
    'tensorboard --logdir {} --host 0.0.0.0 --port 6006 &'
    .format(LOG_DIR)
)
# Tunnel port 6006 (TensorBoard assumed running)
get_ipython().system_raw('lt --port 6006 >> url.txt 2>&1 &')
# Get url
! cat url.txt

最後のcatコマンドで表示されるURLをブラウザで開いてください。
なお、同様の手順でChainer UIなどWebサービス化された他の可視化ツールも動かすことができると思います。

深層学習コンパイラスタック Apache TVM

Apache TVMは、CPU（x86/ARM）やGPU（CUDA/ROCm/OpenGL/OpenCL/Metal）など様々なバックエンドをサポートし、時にはnVidia cuDNNやGoogle TensorRTのパフォーマンスをも凌ぐ（下図参照）深層学習コンパイラです。

Colaboratoryへの導入

TVMのビルドにはかなり時間がかかる上に、多くのライブラリに依存します。そこで、下記サイトのColab Notebookで使われている、ビルド済みバイナリからインストールする方法をお勧めします。

TVM For Fun and Profit Tutorial, at FCRC 2019

C++のビルドが走るため、数十分ほど待つ必要がありますが、以下のコードで、LLVM/CUDA/cuDNN/OpenGLバックエンドに対応したTVMのPythonモジュールがビルドできます。

try:
  import google.colab
  IN_COLAB = True
except:
  IN_COLAB = False

if IN_COLAB:
    ! gsutil cp "gs://tvm-fcrc-binaries-7f775516ff9dfab922c304049f294cec/tvm.tar.gz" /tmp/tvm.tar.gz
    ! mkdir -p /tvm
    ! tar -xf /tmp/tvm.tar.gz --strip-components=4 --directory /tvm
    ! ls -la /tvm
    ! pip install mxnet
    ! bash /tvm/package.sh
    # Add TVM to the Python path.
    import sys
    sys.path.append('/tvm/python')
    sys.path.append('/tvm/topi/python')
    sys.path.append('/tvm/nnvm/python')
    sys.path.append('/tvm/vta/python')
else:
    print("Notebook executing locally, skipping Colab setup ...")

また、下記リンク先の記事に掲載されたコードが、Colaboratory上で実行できることを確認しています。今回は、このコードにカーネルコードの出力コマンドを追加してみました。

AOCL Backend Example

ソースコード

import tvm 
import numpy as np

n = tvm.var("n")
A = tvm.placeholder((n,), name="A")
B = tvm.placeholder((n,), name="B")
C = tvm.compute(A.shape, lambda i: A[i] + B[i], name="C")

s = tvm.create_schedule(C.op)
px, x = s[C].split(C.op.axis[0], nparts=1)

s[C].bind(px, tvm.thread_axis("pipeline"))

fadd_cuda = tvm.build(s, [A, B, C], "cuda", name="myadd")
print(fadd_cuda.imported_modules[0].get_source())

CUDA 出力例

以下の通り、CUDAのカーネルコードが出力されているのが確認できます。

extern "C" __global__ void myadd_kernel0( float* __restrict__ C,  float* __restrict__ A,  float* __restrict__ B, int n) {
  for (int i_inner = 0; i_inner < n; ++i_inner) {
    C[i_inner] = (A[i_inner] + B[i_inner]);
  }
}

OpenCL 出力例

build関数引数のターゲットを変更することで、他のバックエンドのコードも出力可能です。

__kernel void myadd_kernel0(__global float* restrict C, __global float* restrict A, __global float* restrict B, int n) {
  for (int i_inner = 0; i_inner < n; ++i_inner) {
    C[i_inner] = (A[i_inner] + B[i_inner]);
  }
}

Metal 出力例

#include <metal_stdlib>
using namespace metal;

union __TVMArgUnion {
 int v_int;
};

struct myadd_kernel0_args_t {
  int n;
};

kernel void myadd_kernel0(
  device float* C [[ buffer(0) ]],
  device float* A [[ buffer(1) ]],
  device float* B [[ buffer(2) ]],
  constant myadd_kernel0_args_t& arg [[ buffer(3) ]],
) {
  for (int i_inner = 0; i_inner < arg.n; ++i_inner) {
    C[i_inner] = (A[i_inner] + B[i_inner]);
  }
}

エッジ側で深層学習の推論パフォーマンスを極限まで追求するなら、Apache TVMは是非おさえておきたいですね。

まとめ

いかがでしょうか？このように、ColaboratoryならばGPU環境でDeep Learningのコードを手軽に実行＆共有＆再現することができます。個人的には、この環境を活用して最速の速度を競う、チューニング大会を開けたらなぁと思います。