Method replace_batchnorm_to_conv

ppq/IR/morph.py:83–116

Replace Batchnorm to 1D Convolution.

(self, dimension: int = 2)

        ]

    def replace_batchnorm_to_conv(self, dimension: int = 2):
        """ Replace Batchnorm to 1D Convolution. """
        for op in self.graph.operations.values():
            if op.type == 'BatchNormalization':
                ppq_warning(f'Isolated BatchNormalization({op.name}) was detected, '
                            f'PPQ will replace it to 1*1 Convolution({dimension}D).')

                assert len(op.parameters) == 4, "BatchNorm should have 4 parameters, namely alpha, beta, mean, var"
                alpha = op.parameters[0].value
                beta  = op.parameters[1].value
                mean  = op.parameters[2].value
                var   = op.parameters[3].value
                epsilon = op.attributes.get("epsilon", 1e-5)
                
                with torch.no_grad():
                    w = alpha / torch.sqrt(var + epsilon)
                    w = w.reshape([-1, 1] + [1] * dimension)
                    b = alpha * (-mean) / torch.sqrt(var + epsilon) + beta

                op.type = 'Conv'
                op.attributes.clear()
                op.attributes['kernel_shape'] = [1] * dimension
                op.attributes['strides']      = [1] * dimension
                op.attributes['dilations']    = [1] * dimension
                op.attributes['pads']         = [0, 0] * dimension
                op.attributes['group']        = w.numel()

                # remove last 2 variable, make conv has exact 3 input
                self.graph.remove_variable(op.inputs[-1])
                self.graph.remove_variable(op.inputs[-1])

                with torch.no_grad():
                    op.inputs[1].value = w
                    op.inputs[2].value = b

    def replace_batchnorm_to_scale(self, dimension: int = 4):
        """ Replace Batchnorm to Mul + Add.