前回の記事では、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も使ってみたいですね。