Class GraphMerger

ppq/IR/morph.py:501–1075

Graph Merger implements all graph fusion related functions.

class GraphMerger(GraphCommandProcessor):
    """Graph Merger implements all graph fusion related functions."""
    def _acceptable_command_types(self) -> List[GraphCommandType]:
        return [
            # add more extensions in the future
            GraphCommandType.FUSE_BN,
            GraphCommandType.FUSE_BIAS_ADD
        ]

    def process(self, command: GraphCommand) -> Any:
        if command.command_type == GraphCommandType.FUSE_BN:
            return self.fuse_bn()
        if command.command_type == GraphCommandType.FUSE_BIAS_ADD:
            return self.fuse_bias_add()


    def fuse_bn(self):
        search_engine = SearchableGraph(graph=self.graph)
        paths = search_engine.path_matching(
            sp_expr=lambda x: x.type in {'Conv', 'Gemm', 'ConvTranspose'},
            rp_expr=lambda x, y: False,
            ep_expr=lambda x: x.type == 'BatchNormalization',
            direction='down')

        for path in paths:
            path = path.tolist()
            assert len(path) == 2, ('Oops seems we got something unexpected.')

            computing_op, bn_op = path
            assert isinstance(computing_op, Operation) and isinstance(bn_op, Operation)

            if (len(self.graph.get_downstream_operations(computing_op)) != 1 or
                len(self.graph.get_upstream_operations(bn_op)) != 1):
                ppq_warning(f'PPQ can not merge operation {computing_op.name} and {bn_op.name}, '
                            'this is not suppose to happen with your network, '
                            'network with batchnorm inside might not be able to quantize and deploy.')
                continue

            assert len(bn_op.parameters) == 4, 'BatchNorm should have 4 parameters, namely alpha, beta, mean, var'
            alpha = bn_op.parameters[0].value
            beta  = bn_op.parameters[1].value
            mean  = bn_op.parameters[2].value
            var   = bn_op.parameters[3].value
            epsilon = bn_op.attributes.get('epsilon', 1e-5)

            if computing_op.num_of_parameter == 1:
                w = computing_op.parameters[0].value  # no bias.
                assert isinstance(w, torch.Tensor), 'values of parameters are assumed as torch Tensor'
                if computing_op.type == 'ConvTranspose':
                    b = torch.zeros(w.shape[1] * computing_op.attributes.get('group', 1))
                elif computing_op.type == 'Gemm' and computing_op.attributes.get('transB', 0) == 0:
                    b = torch.zeros(w.shape[1])
                else:
                    b = torch.zeros(w.shape[0])
            else:
                w, b = [var.value for var in computing_op.parameters[: 2]]  # has bias.

            if computing_op.type == 'Conv':