本記事はデータ点ごとに勾配を求める手法の中でも、Conv2Dの勾配関数を上書きして求める手法について詳細を記述したものです。
データ点ごとの勾配情報の取り出し
前記事のとおりTensorflowでは、backpropagation時に損失関数$L$に対して、$l$番目のレイヤーの出力$f^l_{N,H,W,C}$で微分したものが伝搬しています。
\delta^l = \Bigl( \frac{\partial L}{\partial f^l_{1,1,1,1}},\frac{\partial L}{\partial f^l_{1,1,1,2}} \dots \frac{\partial L}{\partial f^l_{N,H,W,C}} \Bigr)
よって、各$L_i$に対する微分値は捨てられているのですが、$i$番目のデータ点に対する損失$L_i$が、$j(\neq i)$番目のデータに依存しないとき、$\delta^l$から$\nabla L_i$を復元できます。なぜなら
\frac{\partial L}{\partial f^l_{n,x,y,c}} = \frac{\partial L_n}{\partial f^l_{n,x,y,c}}
となるため、事実上データ点ごとの勾配が伝搬されているためです。$L_i$が$i$番目のデータにのみ依存するという前提は多くのケースで満たすことが可能ですが、学習時のBatch Normalizationでは成り立たないため注意が必要です。
Conv2Dにおける勾配の現状
では、計算グラフ上のどこかのTensorをとってくれば、データ点ごとの勾配$\nabla_{\theta} L_i$が得られるのでしょうか。残念ながらそうではありません。あるレイヤーの出力に関する勾配はデータ点ごとのものが計算グラフ上を流れているのですが、パラメーターに関する勾配を求めるOperatorはバッチ単位の勾配を出力してしまうため、データ点ごとのものはTensorとして計算グラフ上を流れていません。
例えば、Conv2Dのパラメーターに関する勾配を求めるOperatorはconv2d_backprop_filterですが、入力としてデータ点ごとの出力に関する勾配(shape[batch, out_height, out_width, out_channels])をとり、出力としてバッチ単位のパラメーターに関する勾配(shape[filter_height, filter_width, in_channels, out_channels])を返しているのがわかります。
内部的に中間状態としてデータ点ごとのパラメーターに関する勾配(shape[batch, filter_height, filter_width, in_channels, out_channels])を持っていればよいのですが、ぱっと見ではそういうこともなさそうです。ベースとなるであろうcudnnのcudnnConvolutionBackwardFilterがいちいちそのような中間出力を出すとも思えない(未調査)ので、おそらく存在しないでしょう。
というわけで、Conv2Dのデータ点ごとのパラメーターに関する勾配が得たければ自分で実装するしかありません。幸いデータ点ごとの出力に関する勾配は簡単に取得できるため、これを使えば簡単に記述できます。
数式
まずは、データ点ごとの勾配を得るための数式から見ていきます。
Conv2Dの入力を$f$とし、そのサイズを[$H_f$, $W_f$, $C_f$]とします。出力を$g$としそのサイズを[$H_g$, $W_g$, $C_g$]とします。また、その出力に関する勾配を$\delta^{g}$とします。
これらを定義すると、パラメーターに関する勾配$\delta^{\theta}$(shape[$H_{\theta}$,$W_{\theta}$,$C_f$,$C_g$])は、$0 \leqq c_1 < C_f, 0 \leqq c_2 < C_g $の範囲で
\delta^{\theta}_{:,:,c_1,c_2} = \delta^{g}_{:,:,c_2} \ast f_{:,:,c_1}
で求めることができます($\ast$はConv2Dです)。簡単のためstride幅を1とし、paddingは無視しています。このあたりの記事を参考に計算しましたが、間違っていたらごめんなさい。
あとはこれをデータ点ごとに計算するだけです。
独自勾配実装
cudaを書くのは面倒なので出来ればPython上で定義してしまいたいです。上記の計算はチャンネル単位のConvolutionに該当するためdepthwise_conv2dで表現することができます。
depthwise_conv2dのinputとして$f$を、filterとして$\delta^{g}$を与えます。ただし、$f$のチャンネルを[$N$, $H_f$, $W_f$, $C_f$]から[$C_f$, $H_f$, $W_f$, $N$]に並べ替え、$\delta^{g}$のチャンネルを[$N$, $H_g$, $W_g$, $C_g$]から[$H_g$, $W_g$, $N$, $C_g$]に並べ替えます。
そうすると出力は[$C_f$,$H_{\theta}$,$W_{\theta}$,$N * C_g$]の形で得られるため、並び替えることでデータ点ごとの勾配[$N$, $C_f$, $C_g$,$H_{\theta}$,$W_{\theta}$]が得られます。(こちらに実装があります)
これを既存のconv2d_backprop_filterが呼び出されるタイミングで呼び出されるようにすれば、計算グラフ上にデータ点ごとの勾配をあらわすTensorが流れるようになります。
独自勾配実装を呼び出す
Conv2Dの勾配関数を差し替えるとtf.gradientsが呼ばれたタイミングで、独自実装がconv2d_backprop_filterのかわりに計算グラフに組み込まれるようになります。これは、勾配関数を管理しているRegistryをいじることで簡単に達成可能です。
del ops._gradient_registry._registry["Conv2D"]
@ops.RegisterGradient("Conv2D")
def _conv2d_point_backprop_filter(op, out_grad):
input = op.inputs[0]
...
あとは計算グラフ上を流れる勾配が入ったTensorを取得するだけです。このTensorはtf.gradientsの返り値には含まれないため(バッチごとの勾配しか返ってこない)、無理やり計算グラフからとってくる必要があります。
まず、Tensorboardで目的のTensorを出力するOperatorの名前を確認します(TODO:もっとスマートな名前の同定方法を確立したい)。Operator名の末尾に":0"がついたものが、そのOperatorの0番目の出力を意味するTensorです。(tensorflowではOperatorに対して名前がつけられるため、Tensorの名前はOperator名から決まります)
そのTensor名を使って計算グラフからTensorを取得します。あとはsess.runでそのTensorを評価すればデータ点ごとの勾配が手に入ります。
tensor_name = "gradients/layer1/conv0/convolved_grad/per_point_grad:0"
tensor = graph.get_tensor_by_name(tensor_name)
パフォーマンス
上記の方法でforward計算1回、backward計算1回でバッチに所属するすべてのデータ点ごとの勾配を取得することができます。
理論上は速そうなのですが、残念ながら上記実装ではバッチサイズ1として勾配を求める手法よりも遥かに遅いです。これはおそらくdepthwise_conv2dの遅さに由来しています(transposeとか叩いてる方かもしれないけど)。
チューニングしても勝てなさそうなので厳密に裏はとっていないのですが、depthwise_conv2dはcudaによるシンプルな実装がなされているようで、あまり高速化されていません。こちらの記事にもあるように、理論速度との乖離が大きいのではないでしょうか。
cudnnにも直接的に該当する関数はないため(Grouped Convolutionでいける?)、ここを高速化しようと思うとcuda力を高める必要がありそうです。
まとめ
というわけで実装的には面白いのですが、独自勾配関数によるデータ点ごとの勾配の取得は実用的ではありません。
全パラメーターの勾配ではなく低レイヤーのConvolutionの勾配だけ引っ張ってきたいというようなケースでは、既存の計算グラフを最大限使える手法のため他の手法に対する優位性が出てくると思います。
また、depthwise_conv2dが理論値ぐらい速くなればこの実装も有用になると思います。すぐには来ないと思いますが。