[NNabla]構築済みネットワークの中間層間に新たなレイヤーを追加する方法

はじめに

　qiitaへの3回目の投稿です。(article3)
　前回に引き続き、私がnnablaを使っていた中で「こういう情報がqiitaとかにあったらよかったのに」と思いながらなんとか気合いでnnablaのreferenceとdir()(pythonの標準関数。引数のメンバ変数・関数を返してくれる)で見つけてきたことについてまとめます。
#1. 要件
・OS: macOS Catalina (バージョン 10.15.1)
・python: 3.5.4
・nnabla: 1.3.0
#2. ネットワークの構築
　サンプルのネットワークを下記で定義します。(ここまでは前回同様)

article3_add_layer.py

import nnabla as nn
import nnabla.functions as F

# [define network]
x = nn.Variable()
y = F.add_scalar(x, 0.5)  # <-- (1)とおく
y = F.mul_scalar(y, -2)

単純に $y=(x+0.5)\times2$ という形になってます。

#3. 既存の中間層間に新たなレイヤーを追加
前回説明した内容を用いて上記の $y=(x+0.5)\times2$ を$y=(x+0.5)^2\times2$に変える方法を紹介します。コードは下記になります。

article3_add_layer.py

# [get middle variable]
h1 = y.parent.inputs[0]
additional_layer = F.pow_scalar(h1, 2.0)
redefine_layer = F.mul_scalar(additional_layer, **y.parent.info.args)

# [rewire_on]
y.rewire_on(redefine_layer)

動作確認は上記のrewire_onの直前で、printを挟む形で下記で行いました。

article3_add_layer.py

def print_func(f):
    print('{} output = {}'.format(f.name, f.outputs[0].d))

# [print & forward]
x.d.fill(0)
y.forward()
print('--- before ---')
y.visit(print_func)
print('y.d = {}'.format(y.d))
print('')

# [rewire_on]
y.rewire_on(redefine_layer)

# [print & forward]
y.forward()
print('--- after ---')
y.visit(print_func)
print('y.d = {}'.format(y.d))
print('')

出力

--- before ---
AddScalar output = 0.5
MulScalar output = -1.0
y.d = -1.0

--- after ---
AddScalar output = 0.5
PowScalar output = 0.25
MulScalar output = -0.5
y.d = -0.5

###解説

• h1 = y.parent.inputs[0]で(1)の部分を取得しています。
• additional_layer = F.pow_scalar(h1, 2.0)で、(1)の後に新たに追加しようとしていたレイヤーを定義しています。
• この時点でadditional_layerをmul_scalarの出力にrewire_onするとmul_scalar自体も上書きされて消えてしまうので、redefine_layer = F.mul_scalar(additional_layer, **y.parent.info.args)で既存のものと全く同じmul_scalarレイヤーを再定義し、既存のmul_scalarを上書きすることで所望の動作をします。
• y.parent.info.argsは、y.parentの部分でmul_scalarレイヤーを表し、.info.argsでそのレイヤーに与えられた引数を取得しています。つまり、これを使用して既存のmul_scalarレイヤーと全く同じmul_scalarレイヤーを定義できます。
• y.rewire_on(redefine_layer)で、計算グラフ上に再定義されたレイヤーの出力ノードredefine_layerをyで上書きして所望の動作が完了します。
• 最後の動作確認では、rewire_onの前後でレイヤーとして何が接続されているのか？、それぞれの出力は何か？ということを出力しました。レイヤーとして、pow_scalarが増え、数値も数式の通りかと思います。

4. まとめ

　新たなレイヤーを挟み込む方法を紹介しました。これを使用すると、既存の学習済みモデルの各activationの出力に量子化レイヤーを挟み込んだり、Convolution + BatchNormalizationを畳み込んで1つのConvolutionにしたりもできます。次回はこの辺に触れてみようかと思います。