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