Method optimize

ppq/quantization/optim/ssd.py:500–573

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

        self.store_parameter(pair)

    def optimize(
        self,
        graph: BaseGraph,
        dataloader: Iterable,
        executor: BaseGraphExecutor,
        collate_fn: Callable,
        calib_steps: int,
        **kwargs
    ) -> None:
        # restrain maximum img number used for loss checking
        batchsize = 1
        for data in dataloader:
            if collate_fn is not None:
                data = collate_fn(data)
            if isinstance(data, torch.Tensor):
                batchsize = data.shape[0]
            elif isinstance(data, (list, tuple)):
                for value in data:
                    if isinstance(value, torch.Tensor):
                        batchsize = value.shape[0]
                        break
            elif isinstance(data, dict):
                for key, value in data.items():
                    if isinstance(value, torch.Tensor):
                        batchsize = value.shape[0]
                        break
            break

        calib_steps = min(calib_steps, ceil(200 / batchsize))

        all_pairs = self.collect_all_pairs(graph)
        if self.layer_norm:
            self.layer_weight_norm(all_pairs)

        for i in range(self.iteration):
            logger.debug(f'DFQ/SSD Equalization Iteration {i + 1}/{self.iteration}')
            for _,pair in tqdm(enumerate(all_pairs), desc=f'SSD/DFQ Equalization Iteration {i+1}/{self.iteration}', total=len(all_pairs)):
                logger.debug(f"Now Processing Pair {_ + 1}/{len(all_pairs)}: {'--'.join([op.name for op in pair])}")
                self.store_parameter(pair)

                logger.debug(f'Collecting Activation Range for Pair...')
                op_act_range = self.collect_activation_range(pair, executor, dataloader, collate_fn, calib_steps)
                logger.debug(f'Collecting Done!')

                original_weights = self.collect_original_parameter(pair)
                basic_loss = self.test_ssd_loss(pair, executor, dataloader, collate_fn, calib_steps)
                best_loss = basic_loss
                best_idx = -1

                # now apply equalization and estimate loss
                for algo in range(0,4):
                    self.one_step_equalization(pair, op_act_range, algo)
                    self.store_parameter(pair)
                    self.initiate_pair_state(pair)
                    loss = self.test_ssd_loss(pair, executor, dataloader, collate_fn, calib_steps)
                    if algo == 0:
                        logger.debug(f'DFQ Step, Loss Before Equalization {basic_loss} || Loss After Equalization {loss}')
                    else: