こんにちは。R&D改めデータエンジニアリングチームの宮崎です。 普段はTensorFlowを使って、Deep Learningを学習させて画像認識モデルを作ったりしています。
Deep Learningの悩みの一つとして、学習時間が数時間から数日に及び、非常に長いという点が挙げられます。 作業効率やマシンの費用を考えると、少しでも学習時間を短くしたいです。 そこで今回はTensorFlow Profilerを用いてDeep Learningの学習時間を短縮してみたいと思います。
TensorFlow Profilerとは
TensorFlow ProfilerはTensorFlowの計算時に、処理毎の所要時間などパフォーマンス情報を収集し、TensorBoard上で可視化してくれるツールです。 TensorFlow v2.2.0やv2.3.0で機能が強化され、より詳細な分析ができるようになっています。
TensorFlow Profilerのインストール
TensorFlow Profilerを利用するには、通常のTensorFlowの他に、追加プラグインであるtensorboard-plugin-profileをインストールする必要があります。
$ pip3 install tensorflow==2.3.0 $ pip3 install tensorboard==2.3.0 $ pip3 install tensorboard-plugin-profile==2.3.0
MNIST学習の分析
検証にあたって、MNISTというDeep Learningではお馴染みのデータセットをABCIという計算クラウド上で学習させました。ABCIではTesla V100というGPUを利用できます。 なお、残念ながらABCIではTensorFlow Profilerの機能の一つであるメモリプロファイルツールを使うことができませんでした。
TensorFlow Profilerを利用するには、tf.keras.callbacks.TensorBoardの引数profile_batchにプロファイル情報を取得する期間を設定します。
import os import datetime import tensorflow as tf import tensorflow_datasets as tfds from absl import logging def normalize_image(image, label): return tf.cast(image, tf.float32) / 255., label def build_dataset(split, batch_size): dataset, ds_info = tfds.load('mnist', split=split, as_supervised=True, with_info=True) dataset = dataset.map(normalize_image) dataset = dataset.batch(batch_size) if split == 'train': dataset = dataset.repeat() dataset = dataset.shuffle(100) num_examples = ds_info.splits[split].num_examples return dataset, num_examples def build_model(): model = tf.keras.models.Sequential([ tf.keras.layers.Flatten(input_shape=(28, 28, 1)), tf.keras.layers.Dense(128, activation='relu'), tf.keras.layers.Dense(10, activation='softmax') ]) return model def train_with_fit(model, optimizer, loss_fn, epochs, steps_per_epoch, ds_train, ds_test, log_dir, profile_start_step, profile_end_step): model.compile( optimizer=optimizer, loss=loss_fn, metrics=[tf.keras.metrics.SparseCategoricalAccuracy()] ) # プロファイル設定 tboard_callback = tf.keras.callbacks.TensorBoard( log_dir=log_dir, profile_batch=[profile_start_step, profile_end_step]) model.fit(ds_train, epochs=epochs, callbacks=[tboard_callback], validation_data=ds_test, steps_per_epoch=steps_per_epoch) def main(): epochs = 10 batch_size = 128 log_dir = './logs' profile_steps = 20 ds_train, num_train = build_dataset(split='train', batch_size=batch_size) ds_test, num_test = build_dataset(split='test', batch_size=batch_size) steps_per_epoch = num_train // batch_size profile_start_step = int(steps_per_epoch * 1.5) profile_end_step = profile_start_step + profile_steps model = build_model() optimizer = tf.keras.optimizers.Adam() loss_fn = tf.keras.losses.SparseCategoricalCrossentropy() log_dir = os.path.join(log_dir, datetime.datetime.now().strftime('%Y%m%d-%H%M%S')) train_start = datetime.datetime.now() train_with_fit(model, optimizer, loss_fn, epochs, steps_per_epoch, ds_train, ds_test, log_dir, profile_start_step, profile_end_step) train_sec = datetime.datetime.now() - train_start logging.info(f'Train sec: {train_sec}') if __name__ == '__main__': logging.set_verbosity(logging.INFO) main()
コード実行の結果、10エポックの学習には17.732秒かかりました。
TensorFlow Profilerによって指定した学習ステップ間のプロファイル情報が収集され、各処理ごとの計算時間が可視化されます。 これにより、ステップあたりの平均時間が3.4ミリ秒であることがわかりました。 さらに、プロファイルの結果をもとに改善策が示されています。 これらの改善策を試していきたいと思います。
入力パイプラインの最適化
まずは、全体の55%も消費している入力パイプライン最適化のため、build_dataset関数を修正します。
なお、この入力パイプラインの最適化についてはTensorFlow Profilerのチュートリアルに記載の通り、TensorFlow Profilerではお約束の内容のようです。
def build_dataset(split, batch_size): dataset, ds_info = tfds.load('mnist', split=split, as_supervised=True, with_info=True) dataset = dataset.map( normalize_image, num_parallel_calls=tf.data.experimental.AUTOTUNE) # 追記 dataset = dataset.batch(batch_size) dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE) # 追記 if split == 'train': dataset = dataset.repeat() dataset = dataset.shuffle(100) num_examples = ds_info.splits[split].num_examples return dataset, num_examples
実行の結果、学習時間が10.239秒に短縮されました。また、1ステップあたりの平均時間が1.7ミリ秒に減りました。
続いてGPUスレッドの占有化を試してみます。
GPUスレッドの占有化
GPUスレッドの占有化は環境変数をTF_GPU_THREAD_MODE=gpu_privateと設定します。
def main(): os.environ['TF_GPU_THREAD_MODE'] = 'gpu_private' # 追記 (以下略)
実行の結果、学習時間が9.852秒に短縮されました。また、1ステップあたりの平均時間が1.5ミリ秒に減りました。
ここまで順調に学習時間が減ってきました。次は混合精度を用いて、32bit計算を16bit計算に置き換えてみます。
混合精度
混合精度はfloat32の代わりにfloat16を用いて計算することで、メモリの削減を行います。
さらに、今回使用しているV100を含む一部のGPU・TPUで計算速度向上が見込まれます。
from tensorflow.keras.mixed_precision import experimental as mixed_precision #追記 def build_model(): model = tf.keras.models.Sequential([ tf.keras.layers.Flatten(input_shape=(28, 28, 1)), tf.keras.layers.Dense(128, activation='relu'), tf.keras.layers.Dense(10), # activationを外出し tf.keras.layers.Activation('softmax', dtype='float32') # 追記 ]) return model def main(args): os.environ['TF_GPU_THREAD_MODE'] = 'gpu_private' epochs = 10 batch_size = 128 log_dir = './logs' profile_steps = 20 ds_train, num_train = build_dataset(split='train', batch_size=batch_size) ds_test, num_test = build_dataset(split='test', batch_size=batch_size) steps_per_epoch = num_train // batch_size profile_start_step = int(steps_per_epoch * 1.5) profile_end_step = profile_start_step + profile_steps policy = mixed_precision.Policy('mixed_float16') # 追記 mixed_precision.set_policy(policy) # 追記 model = build_model() optimizer = tf.keras.optimizers.Adam() loss_fn = tf.keras.losses.SparseCategoricalCrossentropy() (以下略)
TensorFlow Profiler画面の左下に16bit計算と32bit計算の割合が表示されています。 混合精度を導入した結果、29.6%を16bit計算に置き換えることができました。 しかし残念ながら、学習時間は12.601秒、1ステップあたりの平均時間は2.4ミリ秒にそれぞれ延びてしまいました。
混合精度は効果がみられなかったため無効にし、次はカスタム訓練ループでのTensorFlow Profilerを試してみます。
カスタム訓練ループ
Keras fit関数の代わりに、カスタム訓練ループを用います。
TensorBoardコールバックを設定するだけだったKeras fitと異なり、tf.profiler.experimental.start、tf.profiler.experimental.stop、tf.profiler.experimental.Traceでプロファイルデータの取得を制御する必要があります。
train_with_fitの代わりとなるtrain_with_custom_loopを作成します。
def train_with_custom_loop(model, optimizer, loss_fn, epochs, steps_per_epoch, ds_train, ds_test, log_dir, profile_start_step, profile_end_step): train_loss = tf.keras.metrics.Mean() train_acc = tf.keras.metrics.SparseCategoricalAccuracy() val_loss = tf.keras.metrics.Mean() val_acc = tf.keras.metrics.SparseCategoricalAccuracy() @tf.function def train_step(X, y_true): with tf.GradientTape() as tape: y_pred = model(X) loss = loss_fn(y_true, y_pred) graidents = tape.gradient(loss, model.trainable_weights) optimizer.apply_gradients(zip(graidents, model.trainable_weights)) train_loss.update_state(loss) train_acc.update_state(y_true, y_pred) return loss @tf.function def validation_step(X, y_true): y_pred = model(X) loss = loss_fn(y_true, y_pred) val_loss.update_state(loss) val_acc.update_state(y_true, y_pred) return loss global_step = optimizer.iterations.numpy() summary_writer = tf.summary.create_file_writer(log_dir) train_iter = iter(ds_train) total_steps = epochs * steps_per_epoch logging_interval = math.ceil(steps_per_epoch / 20) with summary_writer.as_default(): for global_step in range(global_step, total_steps): if global_step == profile_start_step: tf.profiler.experimental.start(log_dir) logging.info(f'Start profile at {global_step}') elif global_step == profile_end_step: tf.profiler.experimental.stop() logging.info(f'End profile at {global_step}') with tf.profiler.experimental.Trace('train', step_num=global_step, _r=1): X, y_true = next(train_iter) train_step(X, y_true) if (global_step + 1) % logging_interval == 0: logging.info(f'Steps: {global_step}, ' f'Train Acc: {train_acc.result():.3f}, ' f'Train Loss: {train_loss.result():.3f}') tf.summary.scalar( 'Train/Acc', data=train_acc.result(), step=global_step) tf.summary.scalar( 'Train/Loss', data=train_loss.result(), step=global_step) train_loss.reset_states() train_acc.reset_states() if ((global_step + 1) % steps_per_epoch == 0 or global_step == total_steps - 1): for X, y_true in ds_test: validation_step(X, y_true) logging.info(f'Steps: {global_step}, ' f'Val Acc: {val_acc.result():.3f}, ' f'Val Loss: {val_loss.result():.3f}') tf.summary.scalar( 'Val/Acc', data=val_acc.result(), step=global_step) tf.summary.scalar( 'Val/Loss', data=val_loss.result(), step=global_step) val_loss.reset_states() val_acc.reset_states() def main(): (略) train_start = datetime.datetime.now() # `train_with_fit`から置換 train_with_custom_loop(model, optimizer, loss_fn, epochs, steps_per_epoch, ds_train, ds_test, log_dir, profile_start_step, profile_end_step) train_sec = datetime.datetime.now() - train_start logging.info(f'Train sec: {train_sec}')
実行の結果、学習時間が9.534秒に短縮されました。また、途中スパイクが見られますが、1ステップあたりの平均時間が1.4ミリ秒に減りました。
実験結果まとめ
パターン毎の学習時間をまとめます。一番学習時間が短いのはカスタム訓練ループに入力最適化とGPUスレッド占有化を適用した組み合わせでした。
| 実験 | 学習方法 | 入力 最適化 |
GPU スレッド 占有化 |
混合精度 | 学習時間 | 1ステップ 平均時間 |
検証精度 |
|---|---|---|---|---|---|---|---|
| その1 | Keras fit | ー | ー | ー | 17.732秒 | 3.4ミリ秒 | 97.54% |
| その2 | Keras fit | ● | ー | ー | 10.239秒 | 1.7ミリ秒 | 97.62% |
| その3 | Keras fit | ● | ● | ー | 9.852秒 | 1.5ミリ秒 | 97.80% |
| その4 | Keras fit | ● | ● | ● | 12.601秒 | 2.4ミリ秒 | 97.70% |
| その5 | カスタム 訓練ループ |
● | ● | ー | 9.534秒 | 1.4ミリ秒 | 97.80% |
終わりに
TensorFlow Profilerを用いてMNISTの学習時間の短縮に取り組みました。その結果、見事半分ほどに削減することができました。 MNISTだと学習がすぐに終わってしまうため、効果を感じづらいですが、実データを用いた長時間の学習に適用すればコスト削減に寄与するのではないかと思います。 また、今回はデータ拡張を行っていないなど単純なコードのため、実用時のコードでは、またTensorFlow Profilerの結果も変わってくるのではと思います。 さらに、TensorFlow Profilerには今回用いたサマリ画面以外にも様々なプロファイル画面があり、活用できればより強力なツールになるのではと思います。
ソースコード
今回実験に用いた全体のソースコードは以下となります。
参考
参考にしたサイトです。
