前回の記事では、PyTorchでモデル並列化をして巨大なネットワークの学習法を記事にしましたが、
今回はTensorFlowで巨大なネットワークの学習を行います。
方法としてはUnified Memoryを使う方法もありますが、今回はIBMのTensorFlow-LMS(Large Model Support)を使用する方法を行います。
TensorFlow-LMSはUnified Memoryではなく、データ・スワッピング手法を使っているので高速で出来るようです。(詳しくは論文を参照)
github:https://github.com/IBM/tensorflow-large-model-support
論文:https://arxiv.org/abs/1807.02037
これは、IBMのWMLCE(Watson Machine Learning Community Edition)の一部として利用することが出来ます。
インストール
環境は
Linux(Red Hat Enterprise)
CUDA 9.0
Anaconda
python3.6
tensorflow-gpu 1.10
keras 2.3.1
です。
今回は最新版ではなく、version 1を使用します。
これはgithubでのtflmsv1のブランチに相当します。
インストール方法や使用方法は下記のURLで説明されています。
https://github.com/IBM/tensorflow-large-model-support/tree/tflmsv1
git clone https://github.com/IBM/tensorflow-large-model-support.git -b tflmsv1
pip install ./tensorflow-large-model-support
これでインストール出来ました。
2行目のpip installコマンドが上手くいかない場合は
python ./tensorflow-large-model-support/setup.py install
でもインストール出来ます。
デモを動かす
exampleフォルダ内の3つのサンプルコードを試してみます。
1つ目
python ./tensorflow-large-model-support/examples/mnist_deep_lms.py
2つ目
python ./tensorflow-large-model-support/examples/cnn_mnist_lms.py
3つ目
これはkerasを用いた実装です。--image_sizeの引数は自由な値に出来ます。
python ./tensorflow-large-model-support/examples/Keras_ResNet50.py --image_size 200 --lms
image_sizeを大きめにしても、時間は掛かりますが学習することが出来ました。
これでLarge Model Supportのモードで学習できることが確認できました。
エラー情報
デモを動かしたときに以下のエラーが出て止まってしまうことがあります。
Could not create cudnn handle: CUDNN_STATUS_INTERNAL_ERROR
この場合の対処法は、ターミナルで以下のコマンドで設定を変更するとエラーは無くなります。
export TF_FORCE_GPU_ALLOW_GROWTH=true
参考:https://www.tensorflow.org/guide/gpu https://qiita.com/ysuzuki19/items/b727bbcb45f1cad37630
Kerasでの実装方法
kerasを使用した自分のコードでLarge Model Supportを設定する場合は、以下のコードを冒頭につけるようです。
import tensorflow as tf
from tensorflow.core.protobuf import rewriter_config_pb2
from tensorflow.python.keras import backend as K
config = tf.ConfigProto()
config.graph_options.rewrite_options.memory_optimization = rewriter_config_pb2.RewriterConfig.SCHEDULING_HEURISTICS
config.graph_options.rewrite_options.dependency_optimization = rewriter_config_pb2.RewriterConfig.OFF
K.set_session(tf.Session(config=config))
後は通常のkerasのコードにLMSKerasCallback()を渡すだけのようです。
詳細な説明は開発元のgithubのREADMEにあります。
*追記(2020/5/4)
Large Model Supportで3D U-Netを実装したものをgithubで公開しました。
(Unified Memoryのコードもついでに入っています。)
最新のTensorFlow-LMSを使いたい場合は
今回使ったのはTensorFlowの1.10ですが、最新版は2.xになります。
2.xバージョンでのインストール方法は公式ホームページにあります。
インストールは、channelをIBMのサイトに指定してcondaで出来るようです。-cでチャンネルを指定します。
conda install -c https://public.dhe.ibm.com/ibmdl/export/pub/software/server/ibm-ai/conda/ tensorflow-gpu
使い方も1.xバージョンとは異なるようです。
終わりに
ちなみにIBMのWMLCEは他にもパッケージがあって興味深いです。PowerAIも使ってみたいですね。