ユニファ開発者ブログ

ユニファ株式会社システム開発部メンバーによるブログです。

TensorFlow Profilerを用いた学習時間の短縮

こんにちは。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ミリ秒であることがわかりました。 さらに、プロファイルの結果をもとに改善策が示されています。 これらの改善策を試していきたいと思います。

f:id:unifa_tech:20200907170610p:plain
TensorFlow Profilerの結果

入力パイプラインの最適化

まずは、全体の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ミリ秒に減りました。

f:id:unifa_tech:20200907172211p:plain
TensorFlow Profilerの結果 - 入力パイプライン最適化後

続いてGPUスレッドの占有化を試してみます。

GPUスレッドの占有化

GPUスレッドの占有化は環境変数をTF_GPU_THREAD_MODE=gpu_privateと設定します。

def main():
    os.environ['TF_GPU_THREAD_MODE'] = 'gpu_private' # 追記

    (以下略)

実行の結果、学習時間が9.852秒に短縮されました。また、1ステップあたりの平均時間が1.5ミリ秒に減りました。

f:id:unifa_tech:20200907173415p:plain
TensorFlow Profilerの結果 - GPUスレッド占有化設定後

ここまで順調に学習時間が減ってきました。次は混合精度を用いて、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ミリ秒にそれぞれ延びてしまいました。

f:id:unifa_tech:20200907174900p:plain
TensorFlow Profilerの結果 - 混合精度設定後

混合精度は効果がみられなかったため無効にし、次はカスタム訓練ループでのTensorFlow Profilerを試してみます。

カスタム訓練ループ

Keras fit関数の代わりに、カスタム訓練ループを用います。 TensorBoardコールバックを設定するだけだったKeras fitと異なり、tf.profiler.experimental.starttf.profiler.experimental.stoptf.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ミリ秒に減りました。

f:id:unifa_tech:20200907180450p:plain
TensorFlow Profilerの結果 - カスタム訓練ループ

実験結果まとめ

パターン毎の学習時間をまとめます。一番学習時間が短いのはカスタム訓練ループに入力最適化と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には今回用いたサマリ画面以外にも様々なプロファイル画面があり、活用できればより強力なツールになるのではと思います。

ソースコード

今回実験に用いた全体のソースコードは以下となります。

github.com

参考

参考にしたサイトです。