Help us understand the problem. What is going on with this article?

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

More than 1 year has passed since last update.

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

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

DL framework

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

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のパフォーマンスをも凌ぐ(下図参照)深層学習コンパイラです。

TVM Stack
TVM Backend
TVM nVidia GPU Benchmark

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

stakemura
サービス運用からエッジ組み込みまで、機械学習なんでも屋です。特にCG/CV系のR&D、低レイヤ(C++/GPGPU)、データ分析(Python)、DevOps/MLOps設計の話題が多めです。
Why not register and get more from Qiita?
  1. We will deliver articles that match you
    By following users and tags, you can catch up information on technical fields that you are interested in as a whole
  2. you can read useful information later efficiently
    By "stocking" the articles you like, you can search right away