TensorFlow 2.0のOptimizer

インターンや研究でOptimizerを自作したいときがしばしばあるので（そんなに需要ないかもしれないけど）トライしていきます。

最終目標

TensorFlow 2.xに対応したOptimizerを自作できるようになること

まずは、TensorFlow Core r2.0 におけるOptimizerの基底クラスであるtf.keras.optimizers.Optimizerについて理解していきたいと思います。

以下、公式の和訳とサンプルコード（Google Colabで実行）+コメントです。ちょいちょい直訳っぽいのがあるかもですが、良い訳があればぜひ教えてください。

tf.keras.optimizers.Optimizer の和訳（編集中）

Google ColabでTensorFlow 2.0系使うときは最初にこれ書きます。

from __future__ import absolute_import, division, print_function, unicode_literals
# Install TensorFlow
try:
    # %tensorflow_version only exists in Colab.
    %tensorflow_version 2.x
except Exception:
    pass

import tensorflow as tf

Optimizer クラス

• 最適化のための更新された基底クラス

• このクラスはモデルを訓練するための演算（op, ノード）を追加するためのAPI (Application Programming Interface)を定義する　　

• このクラスを直接利用することはできないが、代わりにtf.keras.optimizers.SGDやtf.keras.optimizers.Adamなどのサブクラスをインスタンス化することができる

使い方

from tensorflow.keras.optimizers import SGD
opt = SGD(lr=0.1)

var1 = tf.keras.backend.variable(1)
var2 = tf.keras.backend.variable(1)
loss = lambda: 3 * var1 * var1 + 2 * var2 * var2
print(loss().numpy())
print(opt.minimize(loss, var_list=[var1, var2]))
print(list(map(lambda x: x.numpy(),  [var1, var2])))
print(loss().numpy())

公式だとvar1とvar2の宣言がなかったので、適当に変数を作って渡しました。

5.0
<tf.Variable 'UnreadVariable' shape=() dtype=int64, numpy=1>
[0.39999998, 0.6]
1.1999999

出力が以上になったので一応最小化できていますね。

Keras モデルと工夫した訓練ループ

• Kerasモデルでは、コンストラクタが呼ばれ(call)たときの代わりに、モデルが初めて呼ばれたときに変数が作成される場合がある
• シーケンシャルモデルやそのサブクラスモデルが入力の形状(shape)を与えられずに定義される場合が例に挙げられる

mnist = tf.keras.datasets.mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

num_hidden = 512
num_classes = 10
tf.keras.backend.set_floatx('float64')

opt = SGD(lr=0.1)
model = tf.keras.Sequential()
model.add(tf.keras.layers.Dense(num_hidden, activation='relu'))
model.add(tf.keras.layers.Dense(num_classes, activation='sigmoid'))
loss_fn = lambda: tf.keras.losses.mse(model(input), output)
var_list_fn = lambda: model.trainable_weights
for input, output in zip(x_train, y_train):
    opt.minimize(loss_fn, var_list_fn)

公式の例に加えて、データセットにMNISTを、隠れ層のノード数と出力数に適当な値を与えたところうまく動かせませんでした...

勾配の前処理

• minimize()を呼んだ場合、勾配の計算と変数への適用を同時に処理する
• 変数への適用をする前に勾配を処理したい場合は代わりにオプティマイザを三段階に分けて使うことができる

1. tf.GradientTapeで勾配を計算する
2. 思うがままに勾配を加工する
3. apply_gradients()で加工された勾配を適用する

opt = SGD(lr=0.1)

with tf.GradientTape() as tape:
    loss = <call_loss_function>
vars = <list_of_variables>
grads = tape.gradient(loss, vars)

processed_grads = [processed_gradient(g) for g in grads]

grad_and_vars = zip(processed_grads, var_list)
capped_grads_and_vars = [(MyCapper(gv[0]), gv[1]) for gv in grads_and_vars]
opt.apply_gradients(capped_grads_and_vars)

具体例挙げるのあきらめました...時間があるときにトライします。
しかし、玄人っぽいオプティマイザの改良ができそうな機構ではありますね！

tf.distribute.Strategyの利用

• このクラスはtf.distribute.Strategyが知られていて、これはすべての複製に渡って勾配が自動で合計されることを意味している
• 勾配の平均を求めるためには、全体(global)のバッチサイズで損失を割る必要があるが、これは訓練または評価ループで構築されたtf.kerasを使えば自動で達成される
• 損失のreductionという要素も確認しよう（平均化する場合はtf.keras.losses.Reduction.SUM_OVER_BATCH_SIZEが、平均化しない場合はtf.keras.losses.Reduction.SUMが設定されているか）
• これらを使わずに勾配の平均を求める場合、tf.math.reduce_sumを使えば、事例ごとの損失を合計して全体のバッチサイズで割ることができる
• tf.distribute.Strategyを使う場合の注意は、tensorのshapeの最初の要素が複製された局所的なバッチサイズであり、1ステップを計算するために使われる複製の数に等しいことである
• その結果、tf.math.reduce_meanを使うと誤った答え（何倍もの大きな勾配）を計算してしまう

変数の制約

(工事中)

スロット

(工事中)

ハイパーパラメータ

(工事中)

カスタマイズしたOptimizerの書き方

(工事中)