Method optimize

ppq/quantization/optim/equalization.py:402–452

(self, graph: BaseGraph, dataloader: Iterable, 
                 executor: BaseGraphExecutor, collate_fn: Callable,
                 **kwargs)

        self.name = 'PPQ Channelwise Split Pass'

    def optimize(self, graph: BaseGraph, dataloader: Iterable, 
                 executor: BaseGraphExecutor, collate_fn: Callable,
                 **kwargs) -> None:

        interested_operations = []
        if self.interested_layers is None:

            for operation in graph.operations.values():
                if operation.type in EQUALIZATION_OPERATION_TYPE:
                    interested_operations.append(operation)
        else:

            for name in self.interested_layers:
                if name in graph.operations:
                    interested_operations.append(graph.operations[name])

        pairs = self.find_equalization_pair(
            graph=graph, interested_operations=interested_operations)

        print(f'{len(pairs)} equalization pair(s) was found, ready to run optimization.')
        for iter_times in tqdm(range(self.iterations), desc='Layerwise Channel Split', total=self.iterations):
            if self.including_act:
                activations = self.collect_activations(
                    graph=graph, executor=executor, dataloader=dataloader, collate_fn=collate_fn,
                    operations=interested_operations)

                for name, act in activations.items():
                    graph.variables[name].value = act # 将激活值写回网络

            for equalization_pair in pairs:
                
                # can not split group convolution.
                is_group_conv = False
                for layer in equalization_pair.downstream_layers + equalization_pair.upstream_layers:
                    if layer.type in {'Conv', 'ConvTranspose'}:
                        group  = layer.attributes.get('group', 1)
                        if group != 1: is_group_conv = True
                
                if is_group_conv: continue
                equalization_pair.channel_split(
                    value_threshold=self.value_threshold,
                    including_bias=self.including_bias,
                    including_act=self.including_act,
                    bias_multiplier=self.bias_multiplier,
                    act_multiplier=self.act_multiplier)

        # channel split progress directly changes fp32 value of weight,
        # store it for following procedure.
        for op in graph.operations.values():
            if isinstance(op, QuantableOperation):
                op.store_parameter_value()