TensorFlowの実装プラクティス
最近忙しくてネタがないので、Qiitaに書いてたこちらの記事を移行してきました。 TensorFlowの実装プラクティス - Qiita
※いろんな媒体に分散して書いちゃってるものを徐々にこちらに一本化していこうかとおもいます
はじめに
TensorFlowを使っていると、たくさんのライブラリや様々な実装の仕方があることがわかると思います。これはtf2が出てきても依然として状況は変わってないように思えます。自由度が高い反面、公開リポジトリのコードを読む際にも、実装の仕方が様々で読み解くのが大変になってしまっていると思います。
今回は、TensorFlowで実験コードを書く際に、個人的に良いと思う実装パターンを書いていこうと思います。結論を言うと、tf.keras.Model
のエコシステムを存分に利用して実装しよう、という主張です。
※ツッコミも大歓迎です。
tf2での実装方法
ここでは主にトレーニングループの実装に焦点をおきます。
トレーニングループの実装
- tf.keras.Model.fit
- custom training loops
- 複雑なトレーニングループを組む必要がある手法の時の選択肢
- GANやSimCLRなど。論文実装はこちらがよく使われている印象
- https://www.tensorflow.org/guide/keras/writing_a_training_loop_from_scratch
- 複雑なトレーニングループを組む必要がある手法の時の選択肢
- tensorflow estimator
- 実験の様々な処理をカプセル化する
- 自由度が高いが、tf.kerasと比べて複雑で可読性が悪くなる印象
- どんどん使われなくなってきていると思うが、ネット上では見かける実装方法
- https://www.tensorflow.org/guide/estimator
主に分けてこれらの実装がよくみられます。特に「こう実装するべき!」という主張はないため、これらが混在している印象です。 tensorflowの初学者は特に実装の方法がたくさんあって苦戦を強いられる状況になってしまっていると思います。 自分自身、TensorFlowでコードを書く時の推奨の実装パターンのようなものがあったら嬉しいなと思っていたので、今回は個人的に良いと思う実装パターンを書いていきたいと思います。
tf.keras.Model
に寄せた実装
ここが今回のメイン部分になります。まずは、tf.keras.Model
に寄せて実装すると何が嬉しいか書いていきます。
いい点
- 便利なコールバックを簡単に使える
- LRスケジュールとか、TensorBoardのログ記録とかも優秀
- compileが優秀
model.save
する際にcompileで設定した情報も保存してくれる- compileしていれば学習途中のoptimizerの状態保存までしてくれる => 実験の再現性
- TensorBoardコールバックなどでも、ここで設定したロスやメトリクスを自動で記録してくれる
逆に、これらの部分は、custom training loops
など、他の方法で実装する際には自前で実装しなければなりません。
tf.keras.Model.fit
を使わない場合
tf.keras.Model
のエコシステムが勝手にやってくれる部分を自前で実装していく必要があります。
tensorboard
例えば、tensorboard用のmetricsやログを取りたかったら、summary_writerを必要な分定義し、トレーニングループ内などでログ記録用のコードを色々と書かなければならなくなり、見栄えが悪くなります。
train_summary_writer = tf.summary.create_file_writer(train_log_dir) test_summary_writer = tf.summary.create_file_writer(test_log_dir) for epoch in range(EPOCHS): for (x_train, y_train) in train_dataset: train_step(model, optimizer, x_train, y_train) with train_summary_writer.as_default(): tf.summary.scalar('loss', train_loss.result(), step=epoch) tf.summary.scalar('accuracy', train_accuracy.result(), step=epoch) for (x_test, y_test) in test_dataset: test_step(model, x_test, y_test) with test_summary_writer.as_default(): tf.summary.scalar('loss', test_loss.result(), step=epoch) tf.summary.scalar('accuracy', test_accuracy.result(), step=epoch) template = 'Epoch {}, Loss: {}, Accuracy: {}, Test Loss: {}, Test Accuracy: {}' print (template.format(epoch+1, train_loss.result(), train_accuracy.result()*100, test_loss.result(), test_accuracy.result()*100)) # Reset metrics every epoch train_loss.reset_states() test_loss.reset_states() train_accuracy.reset_states() test_accuracy.reset_states()
tf.keras.Model.fit
のコールバックであればその設定をするだけでOKです。
tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=log_dir, histogram_freq=1)
optimizer
また、トレーニングは途中で停止・異常終了された時に適切に復旧できるように実装されている必要もあります(トレーニングインスタンスのプリエンプトが主な理由)。 checkpointを保存している例は見ますが、optimizerの状態保存などはあまり考慮されてないケースをよく見かます。これもトレーニングの継続には大事な部分でありますが、自前で実装するのは少し面倒な部分です。
これも、tf.keras.Model.compile
でコンパイルしていれば、(公式で実装済みのoptimizerであれば)model.save
時にデフォルトで保存してくれるようになっています。model.save
のinclude_optimizer
引数で制御することも可能です。
これらのように、tf.keras.Model
に寄せて実装することで、自分で書かなければいけない部分を大幅に減らし、コードの可読性を保つことができます。
ここではひとまずこれらの例を上げましたが、他にも便利な部分や、今後追加されていく機能も多いと思います。
カスタマイズ性
色々と便利な点を書いていきましたが、カスタマイズ性の観点が肝になってくると思います。
そもそも複雑なトレーニングループを書きたいからcustom training roops
で実装しているという方が多いかと思いますが、tf.keras.Model.fit
でトレーニングを実行する場合にも、実はcustom training roops
と同様にトレーニングループを書くことは可能です。
tf.keras.Model.train_stepのオーバーライド
上記公式ドキュメントにも書かれていることなので、知っている人は知ってると思います。当たり前な話ですが、fit
内部で使ってる関数をオーバーライドするということです。
手順
以下のようにtf.keras.Model
のサブクラスを作成します。
class MyModel(tf.keras.Model): """Example in overridden `tf.keras.Model.train_step` Arguments: data: A tuple of the form `(x,)`, `(x, y)`, or `(x, y, sample_weight)`. Returns: The unpacked tuple, with `None`s for `y` and `sample_weight` if they are not provided. """ def train_step(self, data): # If `sample_weight` is not provided, all samples will be weighted # equally. x, y, sample_weight = tf.keras.utils.unpack_x_y_sample_weight(data) with tf.GradientTape() as tape: y_pred = self(x, training=True) loss = self.compiled_loss(y, y_pred) gradients = tape.gradient(loss, self.trainable_variables) self.optimizer.apply_gradients(zip(gradients, self.trainable_variables)) self.compiled_metrics.update_state(y, y_pred, sample_weight) return {m.name: m.result() for m in self.metrics}
参照元。一部改変。
モデルを作成する時に以下のようにカスタムのモデルクラスでラップすればOKです。
def build_model(input_shape: List, num_classes: int): """トレーニングに使用するモデルを作成する. Args: input_shape {List} -- 入力データのshape. num_classes {int} -- クラス数. """ # 例 inputs = tf.keras.Input(shape=input_shape) outputs = tf.keras.layers.Dense(num_classes, activation="softmax")(inputs) model = MyModel(inputs, outputs) # カスタムモデルクラスをしようしない場合は以下 #model = tf.keras.Model(inputs, outputs) return model
たったこれだけでtf.keras.Model.fit
を使ってトレーニングの実行が可能になります。実際、custom training roops
の方法で実装した関数をほぼそのままtf.keras.Model.train_step
に移植するだけで良いと思います。
その他、損失関数やOptimizerなども当然カスタムも可能です。tf.keras.Model.compile
で指定すれば、self.optimizer
やself.compiled_loss
などでトレーニングループ内からアクセスできます。
def custom_loss_func(y: Tensor, y_pred: Tensor) -> Tensor: """カスタムの損失関数を実装する. Args: y {Tensor} -- 例えば教師ラベル y_pred {Tensor} -- 例えばモデルの予測値 Returns: Tensor -- 損失の計算結果 """ loss = y - y_pred return loss model.compile(loss=custom_loss_func, optimizer=custom_optimizer, metrics=[custom_metrics])
tf.keras.Model.fitに寄せた実装にするもう一つの利点は、コードの共通化が可能なところです。ここまで書いてきたカスタマイズできる部分は必要に応じてカスタマイズし、残りは毎回同様のコードを使いまわせることになります(例えば以下)。
def main(argv): if len(argv) > 1: raise app.UsageError('Too many command-line arguments.') # tf.distribute.Strategyを使うかどうか if FLAGS.use_tpu: # Setup tpu strategy cluster = tf.distribute.cluster_resolver.TPUClusterResolver() tf.config.experimental_connect_to_cluster(cluster) tf.tpu.experimental.initialize_tpu_system(cluster) distribute_strategy = tf.distribute.TPUStrategy(cluster) with distribute_strategy.scope(): model = build_model(FLAGS.input_shape, num_classes=FLAGS.num_classes) optimizer = tf.keras.optimizers.Adam(learning_rate=FLAGS.learning_rate) elif FLAGS.use_gpu: # Setup mirrored strategy distribute_strategy = tf.distribute.MirroredStrategy() with distribute_strategy.scope(): model = build_model(FLAGS.input_shape, num_classes=FLAGS.num_classes) optimizer = tf.keras.optimizers.Adam(learning_rate=FLAGS.learning_rate) else: model = build_model(FLAGS.input_shape, num_classes=FLAGS.num_classes) optimizer = tf.keras.optimizers.Adam(learning_rate=FLAGS.learning_rate) model.compile(loss=custom_loss_func, optimizer=optimizer, metrics=["accuracy"]) model.summary() tboard_callback = tf.keras.callbacks.TensorBoard(log_dir=f"{FLAGS.job_dir}/logs", histogram_freq=1) callbacks = [tboard_callback] train_ds = get_dataset(FLAGS.dataset, FLAGS.global_batch_size, "train") valid_ds = get_dataset(FLAGS.dataset, FLAGS.global_batch_size, "valid") for epoch in range(FLAGS.epochs): model.fit(train_ds, validation_data=valid_ds, callbacks=callbacks, initial_epoch=epoch, epochs=epoch+1) model.save(f"{FLAGS.job_dir}/checkpoints/{epoch+1}", include_optimizer=True) model.save(f"{FLAGS.job_dir}/saved_model", include_optimizer=False)
分散学習の場合も、基本的には、ほぼ変わらない実装で機能します。
サンプルリポジトリ
雛形のコードとサンプル実装を以下のリポジトリに載せています。
https://github.com/tonouchi510/tensorflow-design
(今はSimCLRの実装例しか載せてませんが)複雑な手法の実装も可能であることがわかると思います。
なお、宣伝的になってしまいますが、SimCLRのこの実装に関しては、技術書典10でmixi tech note #5
の2章でも掲載予定です。より詳しい情報や、興味がある方は読んでいただけると幸いです。
=> SimCLRの実装としてはMinimalな実装で扱いやすく、分散学習にも対応しているという点で、ある程度需要があるんじゃないかというのもあってこの題材をテーマにしてます。
まとめ
このように、tf.keras.Model.fit
はcompile
で指定できるトレーニング手法しか使えないわけではなく、かなり拡張性が確保されています。tf.keras.Model
のエコシステムを理解すれば、その恩恵を受けつつ、かなり自由度高くトレーニングループを書くことが可能です。
また、自分でコードを書く部分は最小限に抑えられるので、可読性や拡張性の観点で優れているのではないかと思っています。
公式ドキュメントでは、やり方は小さく書かれていますが、あまりこういう実装がいいという主張はなかったように思うので、ここで紹介させていただきました。
現状ではまだ様々な実装のされ方がしていて読むのが辛い状況ですが、これに限らず実装の仕方がもう少し統一されるようになってくれれば良いなぁと思っています。