Edited at

深層学習コンパイラ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から再現することができますが、1からノートを作る場合は以下の手順を踏んでください。


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


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

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


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