[PYTHON] [NNabla] How to add a quantization layer to the middle layer of a trained model

Introduction

This is the 4th post to qiita. (article4) Continuing from last time, while I was using nnabla, I managed to feel like "I wish I had this kind of information in qiita" Summary of what I found in nnabla reference and dir () (a standard python function that returns member variables and functions of arguments) I will.

1. Requirements

· OS: macOS Catalina (version 10.15.1) ・ Python: 3.5.4 ・ Nnabla: 1.3.0

2. Prepare the trained network

This time, from Nnabla trained model, use MobileNet_v1 as follows.

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\'']

Commentary

• nnabla.models provides trained models as a network of nnabla. This time we will use mobilenet_v1 in it.
• Download the trained model of mobilenet_v1 with mobilenet = MobileNet ().
• nnp = mobilenet.nnp stores mobilenet in the variable nnp as a read of nnabla's own file format ".nnp".
• Looking at the list of networks stored in the variable nnp, it says['Validation','Validation5','Runtime','Training'](nnp.get_network_names () This time, we will use Runtime from it as an example. This is the part of net = nnp.get_network ('Validation', batch_size = 1).
• I got the final output y' byy = net.outputs ['y \'']. This will be used later to check the operation. It has nothing to do with what you are trying to do this time.

3. Add a quantization layer to the trained model

Add a quantization layer to each activation, pooling, and affine (identification output) of mobilenet_v1 obtained above. The actual code is below.

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()

The operation check was done below.

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('')

Since the output is long, some parts will be omitted, but it will be in the following form. As a change, MinMaxQuantize (quantization layer) is added after ReLU, ʻAveragePooling, and ʻAffine.

----- before -----
ImageAugmentation
MulScalar
AddScalar
Convolution
BatchNormalization
ReLU
DepthwiseConvolution
BatchNormalization
ReLU
Convolution
...(The following is omitted)...

----- after -----
ImageAugmentation
MulScalar
AddScalar
Convolution
BatchNormalization
ReLU
MinMaxQuantize
DepthwiseConvolution
BatchNormalization
ReLU
MinMaxQuantize
Convolution
...(The following is omitted)...

Commentary

• The top of the process to add a quantization layer is ʻAddQuantizeLayer.add_quantize_layer_all. To give an overview of the process, here we check each variable in the network, and if it is the variable to which the quantization layer should be added (ReLU, ʻAveragePooling, ʻAffine output), ʻAddQuantizeLayer.add_quantize_layer_one Add a quantization layer by.
• self.net.variables is the dict of all network variables. ʻAddQuantizeLayer.add_quantize_layer_alluses this to check each variable in the network. Here,var.parent gets the layer that output the variable, and func.info.type_name` gets the name of that layer.
• ʻAddQuantizeLayer.add_quantize_layer_one uses var_out_new = self.redefine_layer (var_out_cur.parent, q_out) to quantize and redefine the function after that. If you use rewire_on, one variable will be overwritten and disappear, so it is a countermeasure. Specifically, for the part where the layer connection is ...-> ReLU-> Convolution-> ..., branch and connect with ...-> ReLU-> MinMax Quantize. Then you can't suddenly connect like ...-> MinMaxQuantize-> Convolution-> ...(becauserewire_onrequires a variable to be overwritten). At this time, create the same layer as the existingConvolution after MinMaxQuantize, set ...-> MinMaxQuantize-> Convolution (new), and then change the existing Convolution to Convolution (new) . Finally, by overwriting with , the connection...-> MinMaxQuantize-> Convolution-> ...is realized. (I'd like to explain it in the figure.)
• ʻAddQuantizeLayer.redefine_layerdefines the exact same layer as the existing one to handle the above problem._input is the variable after quantization (nn.Variable), _func.inputs [1:] is the other input variable (nn.Variable) (this time the coefficient such as conv), _func. info.args` is the other argument (stride, padding, etc.).

4. Summary

Using the content posted up to Last time, I introduced how to add a quantization layer to an existing network. It is undecided what to post next time.