はじめに

　qiitaへの4回目の投稿です。(article4)
　前回に引き続き、私がnnablaを使っていた中で「こういう情報がqiitaとかにあったらよかったのに」と思いながらなんとか気合いでnnablaのreferenceとdir()(pythonの標準関数。引数のメンバ変数・関数を返してくれる)で見つけてきたことについてまとめます。
#1. 要件
・OS: macOS Catalina (バージョン 10.15.1)
・python: 3.5.4
・nnabla: 1.3.0
#2. 学習済みネットワークの準備
　今回はnnablaの学習済みモデルの中から、下記のようにして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.まとめ
　前回までに投稿した内容を用いて、既存のネットワークに量子化レイヤーを追加する方法を紹介しました。次回は何を投稿しようか未定です。