深層学習コンパイラ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のベンチマークを取ってみた](https://blog.shikoan.com/colab-tpu-resnet-benchmarks/)

## 事前準備

下記ベンチマークで検証した内容は、リンク先の[Colaboratory Note](https://drive.google.com/file/d/1ncuK08ZgOrQZmoMDJ_AjobUpbUV7SMtu/view?usp=sharing)から再現することができますが、１からノートを作る場合は以下の手順を踏んでください。

### Colaboratoryを開く
https://colab.research.google.com/
### 新規ノートブックを作成
`ノートブックを新規作成` > `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
~~~
 
## 利用可能な深層学習フレームワーク

![DL framework](https://cdn-images-1.medium.com/max/1000/1*7C46N8tlMQg28Y631HUZDQ.jpeg)

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

### 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](
https://habrahabr.info/administration/cloud-computing/2397-build-caffe-in-google-colaboratory-a-free-video-card-in-the-cloud.html)

### CNTK

~~~
!pip install cntk
~~~

### Sonnet

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

### Theano

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

## フレームワークのベンチマーク
[簡単なCNNによるディープラーニングライブラリ速度比較](https://qiita.com/daigo0927/items/8092f3ff5276ffc4f088)で紹介されている、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が扱えます→[公式サンプル](https://colab.research.google.com/github/tensorflow/tensorboard/blob/master/docs/r2/tensorboard_in_notebooks.ipynb)** 

Tensorboardのような可視化ツールも、[localtunnel](https://localtunnel.github.io/www/)を使うことで閲覧することができます。
先ほど紹介した、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](https://tvm.ai)は、CPU（x86/ARM）やGPU（CUDA/ROCm/OpenGL/OpenCL/Metal）など様々なバックエンドをサポートし、時にはnVidia cuDNNやGoogle TensorRTのパフォーマンスをも凌ぐ（下図参照）深層学習コンパイラです。

![TVM Stack](https://tvm.ai/images/main/tvm-stack.png)
![TVM Backend](https://github.com/dmlc/web-data/raw/master/tvm/tutorial/tvm_support_list.png)
![TVM nVidia GPU Benchmark](https://tvm.ai/images/autotune-all/nvidia.png)

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

> [TVM For Fun and Profit Tutorial, at FCRC 2019](https://sampl.cs.washington.edu/tvmfcrc/)

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

~~~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](https://github.com/dmlc/tvm/blob/39c8bc2a3da62701f28488920d63fdd2895fe060/docs/deploy/aocl_fpga.md)

### ソースコード

```python
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のカーネルコードが出力されているのが確認できます。

~~~cpp
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関数引数のターゲットを変更することで、他のバックエンドのコードも出力可能です。

~~~cpp
__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 出力例 ###

~~~cpp
#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のコードを手軽に実行＆共有＆再現することができます。個人的には、この環境を活用して最速の速度を競う、チューニング大会を開けたらなぁと思います。