Tensorflow+Estimator+DatasetをMultiGPUで実行する

概要

タイトル通り。Tensorflowに実装されているEstimatorとDatasetを用いて、かつMultiGPUでの学習を試みた際、学習データが非常に大きい場合に
ValueError: Cannot create a tensor proto whose content is larger than 2GB.
が延々と出続ける問題に直面した。
読み込んだ画像をそのままinput_fn内で配列に格納するのではなく、一度TFRecord形式に変換して、TFRecordから読み込むようにすることで解決した。

環境

• Tensorflow 1.14

基礎

[TensorFlow] MirroredStrategyを用いて複数GPU計算を行う
にあるように、Estimatorを用いた複数GPU計算は基本的には容易で、

distribution = tf.distribute.MirroredStrategy()
config = tf.estimator.RunConfig(train_distribute=distribution)
classifier = tf.estimator.Estimator(model_fn=model_fn, config=config)

RunConfigの引数train_distributeに、MirroredStrategy¹を渡して、estimatorの初期化時にconfigに渡せば良い。

ところが上記記事にあるように、

• input_fnの返り値はDataset型でなければならない
• dataset.make_one_shot_iterator().get_next()を返すと input_fn must return a Dataset or a callable.といったエラーが出る

の2点が大きな障害となった。

試したこと

今回、自前で用意した224*224の画像を10000枚ほど学習させることを試みた。

def train_input_fn(train_data, train_label, train_batch_size):
    # train_data:画像をnumpy_arrayに変換したやつ
    # train_label:各画像のラベルを配列にまとめたやつ
    dataset = tf.data.Dataset.from_tensor_slices(({"x": train_data}, train_labels))
    dataset = dataset.shuffle(1000).repeat().batch(train_batch_size)
    return dataset

上記記事ではmnistからDLした画像とラベルの配列をtrain_input_fnに渡している。自前の画像を使う場合には、mnistからデータをダウンロードする代わりに適当に読み出して同じ形式に変換して渡してやればよい。のだが、これを実際に実行すると
ValueError: Cannot create a tensor proto whose content is larger than 2GB.
というエラーが返ってくる。公式ドキュメントでもこの点は指摘されている。

Note that the above code snippet will embed the features and labels arrays in your TensorFlow graph as tf.constant() operations. This works well for a small dataset, but wastes memory---because the contents of the array will be copied multiple times---and can run into the 2GB limit for the tf.GraphDef protocol buffer.

あまり大きなデータに対しては、実行前に定義されるグラフのサイズが大きくなってしまうので、代わりに以下のようにplaceholderを使い、実行時にはIteratorの初期化の際にplaceholderに実際にデータを渡すことが推奨されている。

    features_placeholder = tf.placeholder(train_data.dtype, train_data.shape)
    labels_placeholder = tf.placeholder(train_label.dtype, train_label.shape)

    dataset = tf.data.Dataset.from_tensor_slices((features_placeholder, labels_placeholder))
    # [Other transformations on `dataset`...]
    dataset = ...
    iterator = dataset.make_initializable_iterator()

    sess.run(iterator.initializer, feed_dict={features_placeholder: features,
                                          labels_placeholder: labels})

実際には、Estimatorを使う場合にはSessionが隠蔽されるので上記のように明示的にsess.run(...)とは書けない。そのため、
[Stack over flow] Avoiding tf.data.Dataset.from_tensor_slices with estimator api
などを参考に、一工夫してセッション作成時にイテレータにデータを供給することになる。

ところが先に述べたように、

dataset.make_one_shot_iterator().get_next()を返すと input_fn must return a Dataset or a callable.といったエラーが出る

ので、イテレータを使うことができない。あくまでもinput_fnはDataset型を返す必要があるのである。

解決へ

結論としては、一度読み込む画像を全てTFRecord形式に変換しておき、学習の際にはTFRecordDatasetを用いて読み込むことで、2GBのエラーがでなくなり、無事複数のGPUで実行されることが確認できた。
TFRecordの書き込み・読み込みについては、

• [Qiita]TFRecordの活用例
• TensorFlow推奨フォーマット「TFRecord」の作成と読み込み方法

あたりを参考にした。

def train_input_fn(train_batch_size):
    dataset = tf.data.TFRecordDataset(file_name="hogehoge.tfrecords")
    dataset = dataset.map(parse_feature) # 適当にパースする
    dataset = dataset.shuffle(1000).repeat().batch(train_batch_size)
    return dataset

何がだめだったのか

最初に上げた記事のコードの書き方に沿って自前のでデータを読み込み、from_tensor_slices()で読み込んだ配列をDatasetに入れる場合、当然ながら学習を開始する前に画像の前処理が走るので、静的なグラフ定義に読み込む画像の枚数やサイズといった情報が入ってしまう。そのため大規模なデータを取り扱う際にはグラフサイズが大きくなってしまい、エラーが出る。

一方でTFRecordDataset()を用いる場合には、当然ながらグラフ定義時には「何らかのTFRecordを読み込むらしい」ということしか入らない。実際にデータが供給されるのは、学習を開始してからである。そのためグラフのサイズが大きくなってしまうことを避けられた。

まとめ

TFRecord最高。ありがとう。Tensorflow使うならTFRecord以外は見えない。みんなもEstimatorやDataset、TFRecord形式の恩恵にあずかろう。

1. 上記の記事内ではtf.contrib.distribute.MirroredStrategy()となっているが、TensorFlow1.14の段階ではcontribは取れているようです。 ↩