Add bn fold to quantization-aware training quantizers

[chenbohua3]

At present, many backend frameworks enable the BN-fold function by default, and some of them will adopt per-tensor quantization scheme to quantize the weight tensors. The lack of BN-fold function in quantization-aware training will cause the following two problems:

1. The accuracy gap between the pseudo-quantization (within PyTorch framework) and the real-quantization (within backend framework) would be larger.
2. The loss of accuracy introduced by quantization would be larger.

We are planning to add BN-fold function to quantization-aware training quantizers, especially for QAT_Quantizer and LsqQuantizer. The core idea of bn-fold function is to simulate this folding process in the training graph. The details are described in chapter 3.2 and appendix of the QAT paper

In our implementation, there will be three steps for this function:

1. Find out Conv-Bn pattern groups. we will reuse the TorchModuleGraph and successors modules to find out the groups. To do so, we need to pass a new parameter (e.g. dummy_input) to the __init__ function of the quantizer to get the torchscript graph.

class QAT_Quantizer(Quantizer):
    def __init__(self, model, config_list, optimizer=None, dummy_input=None):
        self._find_conv_bn(model, dummy_input)  # new line
        super().__init__(model, config_list, optimizer)
    
    def _find_conv_bn(self, model, dummy_input, **kwargs):
        self.conv_bn_pattern = {}
        graph = build_module_graph(model, dummy_input)
        for conv in all_convs:
           if relu is matched to conv (using graph):
               self.conv_bn_pattern[conv] = relu

2. Add BN module to QuantizerModuleWrapper, this step may happen when Instantiates QuantizerModuleWrapper or when traverse all the wrappers.

class QuantizerModuleWrapper(torch.nn.Module):
    def __init__(self, module, module_name, module_type, config, quantizer, bn_module):
        ...
        self.bn = bn_module

## or

class QAT_Quantizer(Quantizer):
    def __init__(self, model, config_list, optimizer=None, dummy_input=None):
        .....
        for wrapper in all_module_wrappers:
            if wrapper in fold_list:
                wrapper.bn = bn_module

3. Change the training graph in the forward function of QuantizerModuleWrapper, including two steps:

• If the wrapped module is a folded bn, then we directly return the inputs as the outputs.
• if the wrapped module is a conv with a folded bn, then we add the fold operation.

class QuantizerModuleWrapper(torch.nn.Module):
    ......
    def forward(self, *inputs):
        if self.module is folded bn:
            return inputs
        if 'weight' in self.config['quant_types'] and _check_weight(self.module):
            if self.bn is not None:
                new_weight = fold_operation(self.module.old_weight)
            else:
                new_weight = self.module.old_weight
              
            self.quantizer.quant_grad(
                self.module.old_weight,
                QuantType.QUANT_WEIGHT,
                self, inputs[0])
            result = self.module(*inputs)

This implementation prototype works well for us, looking forward to your suggestions :)

[chenbohua3]

As discussed offline, the current implementation is ok. I will upload a PR about bn-fold. New comments are welcomed.