Method test_ssd_loss

ppq/quantization/optim/ssd.py:431–472

(self,
                    pair: List[Operation],
                    executor: BaseGraphExecutor,
                    data_loader: Iterable,
                    collate_fn: Callable,
                    calib_steps: int
    )

    @torch.no_grad()
    def test_ssd_loss(self,
                    pair: List[Operation],
                    executor: BaseGraphExecutor,
                    data_loader: Iterable,
                    collate_fn: Callable,
                    calib_steps: int
    ) -> float:
        observers = self.build_observer_pair(pair)
        hooks = {op:observers[op].hook for op in observers}
        self.calibrate(pair, data_loader, executor, hooks, collate_fn, calib_steps)
        for _, observer in observers.items():
            observer.render_quantization_config()
        pop_list = []
        for op, observer in observers.items():
            if all([type(var_observer) not in {TorchHistObserver}
                for var_observer in observer._hook._observer_table.values()]):
                    pop_list.append(op)
        for op in pop_list:
            observers.pop(op)
            hooks.pop(op)
        if len(hooks) > 0:
            self.calibrate(pair, data_loader, executor, hooks, collate_fn, calib_steps)
            for _, observer in observers.items():
                observer.render_quantization_config()
        self.calibration_passive_param(pair)
        # calculate loss
        loss = []
        for calib_epoch in range(ceil(calib_steps / len(data_loader))):
            for _,data in enumerate(data_loader):
                if collate_fn is not None:
                    data = collate_fn(data)
                # get the input of first op
                inputs = executor.forward(data, output_names=[pair[0].inputs[0].name])
                # dequant to get fp output
                self.dequantize_pair(pair)
                fp_output = self.run_pair(pair, inputs)
                # restore quant state to get quant output
                self.restore_quantize_state(pair)
                quant_output = self.run_pair(pair, inputs)
                # mse calculation
                loss.append(self.calculate_mse(fp_output, quant_output))
        return torch.stack(loss).mean().item()

    # maintain original parameter for restoration in case of a larger loss after equalization
    def collect_original_parameter(self, pair: List[Operation]) -> Dict[Variable, torch.Tensor]: