benchmark関数を使ってデータセットの処理時間の計測と改善に取り組む

TensorFlowには、tf.data.Dataset APIという入力のパイプラインを実現するための強力な機能があります。入力のパイプラインを最適化することで学習全体を高速化できるため、定量的に計測して改善する価値があります。そこで、本記事ではTensorFlow Datasetsに含まれるbenchmark関数を使ってデータセットの処理時間を測定する方法について紹介します。

ベンチマーク関数

benchmark関数は、tf.data.Datasetのベンチマークを行うための関数です。Datasetを渡すことで、データセットの処理時間に関する以下の情報を出力できます。

総実行時間
セットアップ時間（最初のバッチにかかる時間）
秒あたりに処理できるデータ数（examples/sec）

ここまではドキュメントに書いてある話ですが、実際に実行すると以下のような出力をされて面食らいます。ドキュメントにはstatisticsを返すとさらっと書いてあるだけなので、ソースコードを読んでわかったことを説明しましょう。

{
  "first":{
    "avg":13.807230435075414,
    "duration":0.07242582100025174,
    "num_examples":1
  },
  "first+last":{
    "avg":4597.854424099337,
    "duration":26.09912992699992,
    "num_examples":120000
  },
  "last":{
    "avg":4610.610683215085,
    "duration":26.02670410599967,
    "num_examples":119999
  },
  "raw":{
    "end_time":4150.842121449,
    "first_batch_time":4124.815417343,
    "num_iter":120000,
    "start_time":4124.742991522
  }
}

firstは最初の1バッチに関する統計情報です。最初の1バッチはファイルオープンやネットワーク接続の時間がかかる場合があるため、分けているようです。lastは最初の1バッチ以外を表しています。したがって、first+lastで全バッチを表しています。durationはかかった時間(s)、num_examplesはデータ数、avgは1秒あたりに処理したデータ数を表しています。rawにはパフォーマンスカウンタの値が格納されており、これらの値からdurationを計算しています。

以下では、first+lastのdurationを見て、全体にかかっている時間を計測してみます。

ベンチマーク

benchmark関数を使うためには、TensorFlow Datasetsをインストールする必要があります。この関数は2020年6月25日のPull Requestでマージされているので、TensorFlow Datasets v3.2.0から利用可能です。バージョン3.2.0以上のTensorFlow Datasetsをインストールしましょう。

pip install -U tensorflow-datasets

インストールしたら、必要なパッケージをインポートします。

import tensorflow as tf
import tensorflow_datasets as tfds

まずは、単にデータセットを渡してみましょう。以下のようになりました。

dataset = tfds.load('mnist', split='train', as_supervised=True)
stats = tfds.core.benchmark(
    dataset
    .repeat(2)
)
stats['first+last']['duration']
# 26.867365005000465

次に、mapに処理をさせてみます。少し遅くなることがわかります。

def normalize_img(image, label):
  """Normalizes images: `uint8` -> `float32`."""
  return tf.cast(image, tf.float32) / 255., label

dataset = tfds.load('mnist', split='train', as_supervised=True)
stats = tfds.core.benchmark(
    dataset
    .map(normalize_img)
    .repeat(2)
)
stats['first+last']['duration']
# 39.00072752699998

では、mapを並列実行したらどうでしょうか？先ほどより高速になりました。

dataset = tfds.load('mnist', split='train', as_supervised=True)
stats = tfds.core.benchmark(
    dataset
    .map(
        normalize_img,
        num_parallel_calls=tf.data.experimental.AUTOTUNE
    )
    .repeat(2)
)
stats['first+last']['duration']
# 30.73008285800006

mapの結果をキャッシュしてみましょう。さらに速くなりました。

dataset = tfds.load('mnist', split='train', as_supervised=True)
stats = tfds.core.benchmark(
    dataset
    .map(
        normalize_img,
        num_parallel_calls=tf.data.experimental.AUTOTUNE
    )
    .cache()
    .repeat(2)
)
stats['first+last']['duration']
# 23.48708672599969

mapの前にbatchを入れて、mapをベクトル化します。効果的なようです。

dataset = tfds.load('mnist', split='train', as_supervised=True)
stats = tfds.core.benchmark(
    dataset
    .batch(32)
    .map(
        normalize_img,
        num_parallel_calls=tf.data.experimental.AUTOTUNE
    )
    .cache()
    .repeat(2),
    batch_size=32
)
print(stats['first+last']['duration'])
# 6.9288019289997465

以上で、benchmark関数の使い方の説明は終わりです。benchmark関数を使ってデータセットの処理速度を定量化することで、改善のための一つの指標とすることができます。TensorFlow Profilerなどと合わせて使うことで、ボトルネックを分析し、学習全体を高速化して、生産性の向上やマシンコストの低下を実現させるのに役立ちそうです。

Ahogrammer

Deep Dive Into NLP, ML and Cloud

benchmark関数を使ってデータセットの処理時間の計測と改善に取り組む

ベンチマーク関数

ベンチマーク

参考資料