Tensorflowの学習器が独り占めしている勾配の値をユーザーでも操作したいと思った時、どうすればよいかを調べました。その方法の1つをまとめておきます。

はじめに

import tensorflow as tf
import numpy as np

tensorflowのバージョンは1.8とか1.10くらいです。

導入

\vec{h_1} = \tanh\left(W_1\vec{x} + \vec{b_1}\right) \\
\vec{h_2} = \tanh\left(W_2\vec{h_1} + \vec{b_2}\right) \\
y = \vec{W_3}\cdot\vec{h_2} \\

という計算を考えます。各変数は

$N\times N$行列のパラメータ：$W_1$, $W_2$
$N$次元ベクトルのパラメータ：$\vec{W_3}$, $\vec{b_1}$, $\vec{b_2}$
$N$次元ベクトルの入力・隠れ層：$\vec{x}$, $\vec{h_1}$, $\vec{h_2}$
スカラー値の出力：$y$

という形をしています。

まずは変数を宣言していきます。

ndim = 3
W1 = tf.Variable(np.random.rand(ndim*ndim).reshape(ndim, ndim), dtype=tf.float32)
b1 = tf.Variable(np.zeros(ndim), dtype=tf.float32)
W2 = tf.Variable(np.random.rand(ndim*ndim).reshape(ndim, ndim), dtype=tf.float32)
b2 = tf.Variable(np.zeros(ndim), dtype=tf.float32)
W3 = tf.Variable(np.random.rand(ndim), dtype=tf.float32)

x = tf.placeholder(tf.float32, shape=[ndim])

計算グラフを組み立てます。

h1 = tf.nn.tanh(tf.matmul(W1, tf.reshape(x, (ndim, 1)) + b1))
h2 = tf.nn.tanh(tf.matmul(W2, h1) + b2)
y = tf.reduce_sum(W3*h2)

ちょっと値を与えて計算してみます。

vals = [W1, b1, W2, b2, W3]
init_op = tf.variables_initializer(vals)
with tf.Session() as sess:
    x1 = np.ones(ndim)
    sess.run(init_op)
    ret = sess.run(y, feed_dict=({x: x1}))
    print(ret)

適当な小数が表示されますが、ここでは省略します。

シンプルな学習

$y$のターゲットとなる量はなんでもいいのですが、適当に$|x|$とでもしておきます。
学習するべき変数は

$W_1, W_2, \vec{W_3}, \vec{b_1}, \vec{b_2}$

の5つ。

まず、教師データの置き場所と損失関数の計算グラフを定義します。

t = tf.placeholder(tf.float32, shape=[])
Loss = tf.square(y - t)

optimizerを用意して、学習器を定義します。

optimizer = tf.train.GradientDescentOptimizer(0.001)
train = optimizer.minimize(Loss)

学習を実行し、sess.run(train)前後でW1の内容を出力してみます。

init_op = tf.variables_initializer(vals)
with tf.Session() as sess:
    sess.run(init_op)
    x1 = np.ones(ndim)
    t1 = np.linalg.norm(x1)

    print("Befor")
    print(W1.eval())
    sess.run(train, feed_dict=({x: x1, t: t1}))
    print("After")
    print(W1.eval())

出力

Befor
[[0.69726247 0.9102284  0.79115635]
 [0.35804355 0.3806331  0.09050371]
 [0.6233666  0.23908237 0.3859523 ]]
After
[[0.69727045 0.91023636 0.79116434]
 [0.3581063  0.38069582 0.09056645]
 [0.62341547 0.23913126 0.3860012 ]]

学習により更新されていることが確認できます。
なお、乱数を使っている関係上実行のたびに値が変わるかと思います。

外部から読み込み

ここで、ミニバッチを実装することを考えます。
この程度の例では特徴ベクトルの次元を増やしたらいいだけなのですが、

実装が複雑になってベクトルの次元を増やすだけでも一苦労
特徴ベクトルが可変長でミニバッチとして１つのテンソルにまとめるのが難しい

といった場合があります。

このとき、そもそものデータセットのサブセット($D$)を使ったミニバッチの定義

L(\boldsymbol{x}_D, \boldsymbol{t}_D; \boldsymbol{\omega}) = \sum_{i\in D}l(x_i, t_i; \boldsymbol{\omega})

をもとに、サンプルをいくつか渡してGradientだけ計算しておき、計算グラフの外で集計し、またtrainerに読み込ませることでこのミニバッチを実現してみます。

あ、「損失関数の和の勾配」は「損失関数の勾配の和」と等しいとしています。

\nabla_{\boldsymbol{\omega}}L(\boldsymbol{x}_D, \boldsymbol{t}_D; \boldsymbol{\omega}) = \sum_{i\in D}\nabla_{\boldsymbol{\omega}}l(x_i, t_i; \boldsymbol{\omega})

まず、gradientの数値が入るTensorを取得します。

optimizer = tf.train.GradientDescentOptimizer(0.001)
grads_and_vars = optimizer.compute_gradients(Loss)
grads = [gv[0] for gv in grads_and_vars]

grads_and_varsは2成分のタプルが並んだリストです。
タプルはそれぞれ(gradientのTensor, 対応する変数のTensor)となっています。
そして、これら変数とGradientの組が定まれば変数の更新、つまり学習ができます。

train = optimizer.apply_gradients(grads_and_vars)

今回はgradientのTensorがほしいだけで、特に学習させる変数自体に変更は無いのでそのまま渡します。

早速、適当なサンプル３つでgradientを計算します。

init_op = tf.variables_initializer(vals)
with tf.Session() as sess:
    sess.run(init_op)

    x1 = np.ones(ndim)
    t1 = np.linalg.norm(x1)
    x2 = np.random.rand(ndim)
    t2 = np.linalg.norm(x2)
    x3 = np.random.rand(ndim)
    t3 = np.linalg.norm(x3)

    g1 = sess.run(grads, feed_dict={x: x1, t:t1})
    g2 = sess.run(grads, feed_dict={x: x2, t:t2})
    g3 = sess.run(grads, feed_dict={x: x3, t:t3})

gradsはcompute_gradientsでとってきた勾配のTensorのリストです。
gradsとg1、g2、g3は同じ長さのリストですが、gradにはTensorが、g1〜g3はそのTensorに対応する形状の実際の数値の入ったnumpy配列が並んでいます。

このまま平均を取ります。それで得られた新gradientの値を対応するVariable Tensorに対応付けたdictを作ります。

    gbatch = []
    for g1i, g2i ,g3i in zip(g1, g2, g3):
        gbatch.append((g1i + g2i + g3i)/3.0)
    
    feed_for_train = {}
    for gvar, gval in zip(grads, gbatch):
        feed_for_train[gvar] = gval

このfeed_for_trainを、trainに渡して実行してやれば学習が走ります。

    W1_before = W1.eval()
    sess.run(train, feed_dict=feed_for_train)
    W1_after = W1.eval()

上下のeval()行はちゃんとW1の学習が進んだかのチェックのため、値をとっています。
チェック対象の変数はなんでも良かったのですが、今回は適当にW1を選んでみました。

これまでのコードではサボっていますが、適当なところで値を出力させたら以下のような数字が確認できるかと思います。（乱数のため実行のたびに値自体は変わる）
どこにprintを差し込んだかは記事最後のソースコード全体を参照ください。

出力

1st gradient of W1
[[1.5569324 1.5569324 1.5569324]
 [0.6634847 0.6634847 0.6634847]
 [0.3444534 0.3444534 0.3444534]]
2nd gradient of W1
[[2.7892423  4.2246037  2.9654543 ]
 [1.8092641  2.7403226  1.9235654 ]
 [0.97315514 1.4739467  1.0346347 ]]
3rd gradient of W1
[[2.1652095 2.7425284 2.8009503]
 [1.2192128 1.5442965 1.5771936]
 [0.6687418 0.8470512 0.8650953]]
--------
Averaged gradient of W1
[[2.1704614  2.841355   2.4411123 ]
 [1.2306539  1.6493679  1.3880812 ]
 [0.6621168  0.88848376 0.7480612 ]]
--------
Befor
[[0.11557046 0.92451316 0.43276575]
 [0.7589156  0.34100968 0.4698104 ]
 [0.6307279  0.80514926 0.28046635]]
After
[[0.1134     0.9216718  0.43032464]
 [0.75768495 0.3393603  0.46842232]
 [0.6300658  0.8042608  0.27971828]]
--------
After - Before
[[-0.00217046 -0.00284135 -0.00244111]
 [-0.00123066 -0.00164938 -0.00138807]
 [-0.00066209 -0.00088847 -0.00074807]]

サンプルごとに違うGradientの値が平均され、その結果に0.001倍されたものでW1の値が更新されています。
つまり、Gradientの値の上書きに成功しました。

補足

なお、OptimizerからGradientを取り出し、クリッピングしてもう一度適用するという例が公式ドキュメント
https://www.tensorflow.org/api_docs/python/tf/train/Optimizer
にありますが、これはこの一連の作業が計算グラフに組み込めるので、feed_dictに入力ベクトルと教師信号さえ与えたら完結します。
今回やろうとしたことは、複数の計算グラフ適用を超えて値をやり取りする必要があるので、計算グラフの途中（特にGradient）で値を差し込む方法を改めて確認する必要がありました。

なお、自分の場合はGraph Convolutional Networkで遊んでいました。
グラフはサンプルごとにサイズが違うため、可変長時系列と同じ問題で、単純に次元を増やすミニバッチを使えません。
他のGCNライブラリでは複数のグラフを1つのグラフとし、グラフ間にエッジがないものとして取り扱う方法をとっていました。
僕もそれを試そうとしましたが、悪戦苦闘した挙句失敗しましたので、とりあえずの方法ということでこの記事の手段に行き着きました。

あと、gradientの平均を取る操作もTensorflowの計算グラフに入れられないか、そしたらOpenMP並列、GPU演算使って平均取れるのにとは思いますが、そもそもこんなことをし始める時点でパフォーマンスよりもとにかく動く実装を得たいという場合の処方箋です。

ただサンプルごとの勾配計算するところをプロセス間、ノード間でやりとりできたら分散計算できるんじゃないかなっていう妄想。

その他

設計はちゃんとしましょう（自戒）。

ソースコード全体

import tensorflow as tf
import numpy as np

ndim = 3
W1 = tf.Variable(np.random.rand(ndim*ndim).reshape(ndim, ndim), dtype=tf.float32, name="W1")
b1 = tf.Variable(np.zeros(ndim), dtype=tf.float32, name="b1")
W2 = tf.Variable(np.random.rand(ndim*ndim).reshape(ndim, ndim), dtype=tf.float32, name="W2")
b2 = tf.Variable(np.zeros(ndim), dtype=tf.float32, name="b2")
W3 = tf.Variable(np.random.rand(ndim), dtype=tf.float32, name="W3")

x = tf.placeholder(tf.float32, shape=[ndim])

vals = [W1, b1, W2, b2, W3]

h1 = tf.nn.tanh(tf.matmul(W1, tf.reshape(x, (ndim, 1)) + b1))
h2 = tf.nn.tanh(tf.matmul(W2, h1) + b2)
y = tf.reduce_sum(W3*h2)

t = tf.placeholder(tf.float32, shape=[])
Loss = tf.square(y - t)

optimizer = tf.train.GradientDescentOptimizer(0.001)
grads_and_vars = optimizer.compute_gradients(Loss)
grads = [gv[0] for gv in grads_and_vars]

for gv in grads_and_vars:
    print(gv[1], "-grad->", gv[0])

train = optimizer.apply_gradients(grads_and_vars)


init_op = tf.variables_initializer(vals)
with tf.Session() as sess:
    sess.run(init_op)

    x1 = np.ones(ndim)
    t1 = np.linalg.norm(x1)
    x2 = np.random.rand(ndim)
    t2 = np.linalg.norm(x2)
    x3 = np.random.rand(ndim)
    t3 = np.linalg.norm(x3)

    g1 = sess.run(grads, feed_dict={x: x1, t:t1})
    g2 = sess.run(grads, feed_dict={x: x2, t:t2})
    g3 = sess.run(grads, feed_dict={x: x3, t:t3})

    gbatch = []
    for g1i, g2i ,g3i in zip(g1, g2, g3):
        gbatch.append((g1i + g2i + g3i)/3.0)
    
    feed_for_train = {}
    for gvar, gval in zip(grads, gbatch):
        feed_for_train[gvar] = gval

    print("1st gradient of W1")
    print(g1[0])
    print("2nd gradient of W1")
    print(g2[0])
    print("3rd gradient of W1")
    print(g3[0])
    print("--------")
    print("Averaged gradient of W1")
    print(gbatch[0])
    print("--------")
   
    W1_before = W1.eval()
    print("Befor")
    print(W1_before) 

    sess.run(train, feed_dict=feed_for_train)
    W1_after = W1.eval()
    print("After")
    print(W1_after) 
    print("--------")
    print("After - Before")
    print(W1_after - W1_before)