[NNabla]学習済みモデルの中間層に量子化レイヤーを追加する方法

# はじめに
　qiitaへの4回目の投稿です。(article4)
　[前回](https://qiita.com/__NJ__000/items/c0c24f15bb1253323c0a)に引き続き、私がnnablaを使っていた中で「こういう情報がqiitaとかにあったらよかったのに」と思いながらなんとか気合いで[nnablaのreference](https://nnabla.readthedocs.io/en/latest/)と`dir()`(pythonの標準関数。引数のメンバ変数・関数を返してくれる)で見つけてきたことについてまとめます。
#1. 要件
・OS: macOS Catalina (バージョン 10.15.1)
・python: 3.5.4
・nnabla: 1.3.0
#2. 学習済みネットワークの準備
　今回は[nnablaの学習済みモデル](https://blog.nnabla.org/ja/release/v1-0-20/)の中から、下記のようにしてMobileNet_v1を使用します。

```article4_add_quantization_for_network.py
import nnabla as nn
import nnabla.functions as F
import nnabla.parametric_functions as PF
from nnabla.models.imagenet import MobileNet

if __name__ == "__main__":
    # [load network]
    mobilenet = MobileNet()
    nnp = mobilenet.nnp
    net = nnp.get_network('Runtime', batch_size=1)
    y =  net.outputs['y\'']
```
### 解説
* `nnabla.models`には学習済みモデルがnnablaのネットワークとして用意されています。今回はその中でmobilenet_v1を使用します。
* `mobilenet = MobileNet()`でmobilenet_v1の学習済みモデルをダウンロードします。
* `nnp = mobilenet.nnp`でnnabla独自の「.nnp」というファイル形式を読み込んだ形としてmobilenetが変数`nnp`に保持されます。
* 変数`nnp`に保存されているネットワークのリストを見てみると、`['Validation', 'Validation5', 'Runtime', 'Training']`となっていて(`nnp.get_network_names()`で取得)、今回は例としてその中から`Runtime`を使用します。これが`net = nnp.get_network('Validation', batch_size=1)`の部分です。
* `y =  net.outputs['y\'']`によって、最終出力`y'`を取得しました。これはあとで動作確認のために使用します。今回やろうとしていることに直接的には関係ありません。

# 3. 学習済みモデルに量子化レイヤーを追加
　上記で取得したmobilenet_v1の各activation、pooling、affine(識別用出力)に量子化レイヤーを追加します。実際のコードは下記になります。

```article4_add_quantization_for_network.py
class AddQuantizeLayer:
    def __init__(self, _net):
        self.net = _net

    def add_quantize_layer_all(self):
        # [quantize]
        count = 0
        for key in self.net.variables:
            var = self.net.variables[key]
            func = var.parent
            if type(func) != type(None):
                if func.info.type_name in ['ReLU', 'AveragePooling', 'Affine']:
                    count = self.add_quantize_layer_one(var, count)

    def add_quantize_layer_one(self, _var, _count):
        var_out_cur = _var.function_references[0].outputs[0]
        # [quantize]
        q_out = PF.min_max_quantize(_var, ql_min=0, ql_max=255, x_min_max=True, name='MinMaxQuantize_{}'.format(_count))
        # [redefine function]
        var_out_new = self.redefine_layer(var_out_cur.parent, q_out)
        var_out_cur.rewire_on(var_out_new)
        return _count + 1

    @staticmethod
    def redefine_layer(_func, _input):
        if _func.info.type_name == 'DepthwiseConvolution':
            return F.depthwise_convolution(_input, *_func.inputs[1:], **_func.info.args)
        elif _func.info.type_name == 'Convolution':
            return F.convolution(_input, *_func.inputs[1:], **_func.info.args)
        elif _func.info.type_name == 'AveragePooling':
            return F.average_pooling(_input, **_func.info.args)
        elif _func.info.type_name == 'Affine':
            return F.affine(_input, *_func.inputs[1:], **_func.info.args)
        elif _func.info.type_name == 'Softmax':
            return F.softmax(_input, **_func.info.args)
        else:
            print('[error] redefine_layer()')
            print('_func is unexpected layer.')
            print('_func.info.type_name = {}'.format(_func.info.type_name))
            exit(0)

# [quantize]
AddQuantizeLayer_class = AddQuantizeLayer(net)
AddQuantizeLayer_class.add_quantize_layer_all()
```
動作確認は下記で行いました。

```article4_add_quantization_for_network.py
def print_func(f):
    print('{}'.format(f.name))

print('----- before -----')
y.visit(print_func)
print('')

# [quantize]
AddQuantizeLayer_class = AddQuantizeLayer(net)
AddQuantizeLayer_class.add_quantize_layer_all()

print('----- after -----')
y.visit(print_func)
print('')
```
出力は長いので一部省略しますが下記の形になります。変化としては、`ReLU`、`AveragePooling`、`Affine`の後に`MinMaxQuantize`(量子化レイヤー)が追加されています。

```
----- before -----
ImageAugmentation
MulScalar
AddScalar
Convolution
BatchNormalization
ReLU
DepthwiseConvolution
BatchNormalization
ReLU
Convolution
...(以下略)...

----- after -----
ImageAugmentation
MulScalar
AddScalar
Convolution
BatchNormalization
ReLU
MinMaxQuantize
DepthwiseConvolution
BatchNormalization
ReLU
MinMaxQuantize
Convolution
...(以下略)...
```

###解説
* 量子化レイヤーを追加する処理のtopは`AddQuantizeLayer.add_quantize_layer_all`です。処理の概要を説明すると、ここではネットワークの各変数を確認し、それが量子化レイヤーを追加すべき変数(`ReLU`、`AveragePooling`、`Affine`の出力)であれば`AddQuantizeLayer.add_quantize_layer_one`によって量子化レイヤーの追加をします。
* `self.net.variables`はネットワークの全変数の`dict`です。`AddQuantizeLayer.add_quantize_layer_all`では、これを用いてネットワークの各変数を確認しています。ここで、`var.parent`でその変数を出力したレイヤー、`func.info.type_name`でそのレイヤーの名称を取得しています。
* `AddQuantizeLayer.add_quantize_layer_one`では、`var_out_new = self.redefine_layer(var_out_cur.parent, q_out)`によって量子化とそのあとの関数の再定義をしています。`rewire_on`を使用すると変数が一つ上書きされて消えるのでその対策です。具体的には、レイヤーの接続が`...->ReLU->Convolution->...`となっている箇所に対して、分岐させて`...->ReLU->MinMaxQuantize`と接続してからいきなり`...->MinMaxQuantize->Convolution->...`のように接続させることはできません(`rewire_on`では上書きされる変数が必要になるため)。この時は既存の`Convolution`と同じレイヤーを`MinMaxQuantize`の後に作成し、`...->MinMaxQuantize->Convolution(new)`としてから既に存在していた`Convolution`を`Convolution(new)`で上書きする形で最終的に`...->MinMaxQuantize->Convolution->...`という接続を実現します。(そのうち図での説明にしたい。)
* `AddQuantizeLayer.redefine_layer`では上記の問題を処理するために、既存のものと全く同じレイヤーを定義しています。`_input`が量子化後の変数(`nn.Variable`)、`_func.inputs[1:]`が他の入力変数(`nn.Variable`)(今回はconvなどの係数)、`_func.info.args`がその他の引数(strideやpaddingなど)となっています。

#4.まとめ
　[前回](https://qiita.com/__NJ__000/items/c0c24f15bb1253323c0a)までに投稿した内容を用いて、既存のネットワークに量子化レイヤーを追加する方法を紹介しました。次回は何を投稿しようか未定です。