Method optimize

ppq/quantization/optim/training.py:531–566

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

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

        blocks = self.split_graph_into_blocks(
            graph=graph, executing_order=executor._executing_order,
            blocksize=self.block_size, interested_layers=self.interested_layers)

        # ready for finetuning, print information.
        print('')
        print('Check following parameters:')
        print(f'Interested Layers:         {self.interested_layers}')
        print(f'Num of blocks:             {len(blocks)}')
        print(f'Steps:                     {self.steps}')
        print(f'collecting_device:         {self.collecting_device}')
        print('') # blank line

        # do per-block finetune
        for block_idx, block in enumerate(blocks):
            qt_inputs, fp_outputs = self.collect(
                graph=graph, block=block, executor=executor, 
                dataloader=dataloader, collate_fn=collate_fn, 
                collecting_device=self.collecting_device)

            print(f'# Block [{block_idx + 1} / {len(blocks)}]: '
                  f'[{block.sp.name} -> {block.ep.name}]')
            pre_loss, post_loss = self.correct_bias(
                qt_inputs=qt_inputs, fp_outputs=fp_outputs,
                block=block, executor=executor, graph=graph)
            print(f'# Tuning Finished  : ({pre_loss:.4f} -> {min(pre_loss, post_loss):.4f}) [Block Loss]')
            print('') # blank line


class LearnedStepSizePass(TrainingBasedPass):