TensorFlowの勾配計算の闇

研究でTensorFlowを使ってて勾配計算における闇を発見してそれに2日位悩まされたので記事にしておく
TensorFlow 1.xの話です

・TensorFlowの勾配計算って？ああ！！

通常TensorFlowを扱う時、以下の様に書いて勾配計算って全く気にせずに書くと思います。(というかTensorFlow1.x系列だと勾配自体がブラックボックスに近い技術になってる気がする)

import tensorflow as tf

optimizer = #optimizerの宣言
train_op = optimizer.minimize(#誤差)

RNNなどで勾配クリッピングする際は以下の様に書いて勾配を少し意識すると思います。

import tensorflow as tf

optimizer = # optimizerの宣言
grad, vars = zip(*optimizer.compute_gradients(# 誤差))
cliped_grad = # 勾配クリッピング処理
train_op = optimizer.apply_gradients(zip(cliped_grad, vars))

それ以外では勾配を意識することは無いと思います。ところがどっこい、optimizerのminimizeメソッド内部処理(https://github.com/tensorflow/tensorflow/blob/05ab6a2afa2959410d48aab2336cea9dc1e2c13e/tensorflow/python/training/optimizer.py#L355)を見てみるとcompute_gradientsメソッドを呼び出して勾配を計算し、その勾配をapply_gradientsメソッドで適用してることがわかると思います。minimizeメソッドを使用することでということは知らないうちに勾配を計算して、適用してることになります
これから勾配計算の仕様と闇に触れます。

・勾配計算と適用の仕様

1.勾配計算の仕様

基本的にはcompute_gradientsメソッドに誤差を引数として実行することで勾配が計算されます。
これの戻り値は勾配とその勾配を適用するための変数がペア(勾配と変数のタプル)になったリストになります。なのでZip関数を使用することで勾配と変数をそれぞれリストとして取り出すことができます。型は勾配：tf.Tensor、変数：tf.Variableとなります。

2.勾配適用の仕様

勾配をモデルに適用するするためにはapply_gradientsメソッドに勾配と変数のタプルのリストを引数として実行することで適用されます。
なので一回勾配に対してクリッピングなどの処理を行なった場合は再度Zip関数を使用してペアのリストにする必要があります。
次からその勾配計算に関する闇に触れます

・闇1．勾配クリッピング

1つ目の闇は「勾配クリッピング」です。
TensorFlowには値やノルム等で値をクリップする手段が多数用意されています。しかし、勾配(tf.Tensorのリスト)を扱える関数はtf.clip_by_global_norm以外存在しないです。その他の方法でクリッピングしたい場合はfor文で回さないといけないです・・・・（なんともめんどくさい）
TensorFlow公式チュートリアルもtf.clip_by_global_normでクリッピングをしているので、それでやれってことなんでしょう・・・
以下コードです

import tensorflow as tf

optimizer = # optimizerの宣言
grad, vars = zip(*optimizer.compute_gradients(# 誤差))

# グローバルノルムでのクリッピング
cliped_gnorm_grad, norm = tf.clip_by_global_norm(grad, # クリップする値)

# それ以外でのクリッピング
cliped_other_grad = [# one_gradに対してtf.clip_by_valueなどの関数を適用する for one_grad in grad]

・闇2．勾配を計算グラフ外に取り出す

2つ目の闇は「勾配を計算グラフ外に取り出す」です。
なんでいちいち勾配を計算グラフ外に取り出すの？って思うかもしれないですがモデルが巨大(EfficentNet-B8とか)になってまともにバッチ数を稼げない時に、少ないバッチ数で勾配を数回計算してその勾配の平均値を取ったりすることで擬似的にバッチ数を稼ぐことができます。なのでそういった時にどうしても勾配の値を一回計算グラフ外に出したりします。
普通にTensorFlowを使用している方なら「余裕じゃん、勾配自体をsess.runで呼び出せば出てくるでしょ？(以下のコード)」と思うのですが、そこが闇ポイントです。

import tensorflow as tf

optimizer = # optimizerの宣言
grad, vars = zip(*optimizer.compute_gradients(# 誤差))

init = tf.global_variables_initializer()
with tf.Session() as sess:

init.run()

gradients = sess.run(grad, feed_dict={# データとプレースホルダ})

上記のコードは正常に動作しません(コメントの部分を正しく書き直しても動作しません)、「嘘だ！！」と思う方は実際に書いてみて実行していただければ
エラーの内容としましてはsess.runで実行するオペレーションがNoneになってるよってエラーです。
「tf.Tensorを指定して実行しているのに何故？」「tf.TensorなのにNoneっておかしくね？」って首を傾げる人は多いと思います。このエラーが出る原因は、勾配自体が特殊なTensorになっているからです。勾配のリストの一部を表示した時、tf.Tensorと型が出ますがその後の名前がこのエラーの原因です。名前の後ろに"Control_dependency"とあると思いますこれはTensorFlowの中のTensorでも「実行する前に値が入ってないと行けないTensor」(https://www.tensorflow.org/versions/r1.15/api_docs/python/tf/control_dependencies)になります。なのでsess.runで実行することはできません。
「じゃあどうするの？」ってなりますがsess.runにvarsつまり、勾配をセットする変数を指定すること(以下のコード)で取り出すことが可能になります。

import tensorflow as tf

optimizer = # optimizerの宣言
grad, vars = zip(*optimizer.compute_gradients(# 誤差))

init = tf.global_variables_initializer()
with tf.Session() as sess:

init.run()

gradients = sess.run(vars, feed_dict={# データとプレースホルダ})

勾配計算なのに変数を指定する・・・・闇が深い・・・

・闇3．勾配を計算グラフ外に取り出すオペレーションを文字列で指定する

3つ目の闇は「勾配を計算グラフ外に取り出すオペレーションを文字列で指定する」です。
グラフを外部から読み込んだ時、どうしても文字列でオペレーションを指定する必要があるので時々こうなります・・・
上の2つ目の闇を読んだ方は「変数の名前の文字列のリストを指定すればいい」と考えてコードを打ち込んで実行すると思いますが、エラーもなく実行できますが、それが深い深い闇の入り口です・・・

import tensorflow as tf

optimizer = # optimizerの宣言

grad, vars = zip(*optimizer.compute_gradients(# 誤差))

gradient_op = [var.name for var in vars]
init = tf.global_variables_initializer()

with tf.Session() as sess:

init.run()

gradients = sess.run(gradient_op, feed_dict={# データとプレースホルダ})

上記コードが実行できるのでこの勾配に対して、色々と操作してモデルに勾配適用しようと思ったら何故かエラーが出ます。何故かというときちんと勾配がグラフ外に取り出せてないからです。printとかで計算した勾配を見てみると1組の整数型のタプルを要素に持つndarray(データタイプはオブジェクト)のリスト(正しい勾配はfloat32のデータを持つndarrayのリスト)になっています。一見すると勾配計算間違えたのかな？って思いますがオペレーションの設定ミスです。
原因は変数は変数であり実際の勾配の計算オペレーションとは違うところです。変数の名前はkernelという名前が最後につき、勾配はControl_dependencyという名前がついています。実際のオペレーションはControl_dependencyで行われ、それがkernelに代入されるという流れなので文字列で指定する場合は勾配の方のTensorの名前を設定する必要があります(以下コード)

import tensorflow as tf

optimizer = # optimizerの宣言

grad, vars = zip(*optimizer.compute_gradients(# 誤差))

gradient_op = [one_grad.name for one_grad in grad]
init = tf.global_variables_initializer()

with tf.Session() as sess:

init.run()

gradients = sess.run(gradient_op, feed_dict={# データとプレースホルダ})

同じ勾配計算なのに、文字列だと勾配、テンサーだと変数を指定しないといけない・・・闇が深すぎる・・・・