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計算を行う][multi-gpu-qiita]
にあるように、`Estimator`を用いた複数GPU計算は基本的には容易で、

```py
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`[^1]を渡して、`estimator`の初期化時に`config`に渡せば良い。

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

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

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

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

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

```py
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.`
というエラーが返ってくる。[公式ドキュメント][tensorflow-importingdata]でもこの点は指摘されている。
> 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`に実際にデータを渡すことが推奨されている。

```py
    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][IteratorInitializerHook]
などを参考に、一工夫してセッション作成時にイテレータにデータを供給することになる。

ところが先に述べたように、
> `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の活用例][qiita-tfrecord]
* [TensorFlow推奨フォーマット「TFRecord」の作成と読み込み方法][tensorflow-dataset-api]

あたりを参考にした。

```py
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
```

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

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

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

[IteratorInitializerHook]:https://stackoverflow.com/questions/52266000/avoiding-tf-data-dataset-from-tensor-slices-with-estimator-api
[qiita-tfrecord]:
https://qiita.com/wasnot/items/9b64550237a3c5267bfd

[tensorflow-dataset-api]:https://www.tdi.co.jp/miso/tensorflow-tfrecord-01
[multi-gpu-qiita]:https://qiita.com/shu-yusa/items/e93e934a14849541de78
[tensorflow-importingdata]:https://www.tensorflow.org/guide/datasets#consuming_numpy_arrays