Method collect

ppq/quantization/optim/training.py:224–298

Collect training data for given block. This function will collect fp32 output and quantized input data by executing your graph twice. For collecting fp32 output, all related operations will be dequantized. For collecting quantized input, all related opera

(
        self, graph: BaseGraph, block: TrainableBlock, executor: TorchExecutor, 
        dataloader: Iterable, collate_fn: Callable, collecting_device: str, steps: int = None,
        expire_device: str = 'cpu')

        return ret

    def collect(
        self, graph: BaseGraph, block: TrainableBlock, executor: TorchExecutor, 
        dataloader: Iterable, collate_fn: Callable, collecting_device: str, steps: int = None,
        expire_device: str = 'cpu') -> Tuple[List[Dict[str, torch.Tensor]], List[Dict[str, torch.Tensor]]]:
        """
        Collect training data for given block.
        This function will collect fp32 output and quantized input data by
            executing your graph twice.
        For collecting fp32 output, all related operations will be dequantized.
        For collecting quantized input, all related operations' quantization state will be restored.

        collecting device declares where cache to be stored:
            executor - store cache to executor device.(default)
            cpu      - store cache to system memory.
            cuda     - store cache to gpu memory.(2x speed up)
            disk     - not implemented.

        Args:
            block (TrainableBlock): _description_
            executor (TorchExecutor): _description_
            dataloader (Iterable): _description_
            collate_fn (Callable): _description_
            collecting_device (str): _description_

        Returns:
            _type_: _description_
        """
        def cache_fn(data: torch.Tensor):
            # TODO move this function to ppq.core.IO
            if not isinstance(data, torch.Tensor):
                raise TypeError('Unexpected Type of value, Except network output to be torch.Tensor, '
                                f'however {type(data)} was given.')
            if collecting_device == 'cpu': data = data.cpu()
            if collecting_device == 'cuda': data = data.cuda()
            # TODO restrict collecting device.
            return data

        with torch.no_grad():
            try:
                if len(dataloader) > 1024:
                    ppq_warning('Large finetuning dataset detected(>1024). '
                                'You are suppose to prepare a smaller dataset for finetuning. '
                                'Large dataset might cause system out of memory, '
                                'cause all data are cache in memory.')
            except Exception as e:
                pass # dataloader has no __len__

            quant_graph = QuantableGraph(graph) # helper class
            fp_outputs, qt_inputs = [], []
            
            cur_iter = 0
            # dequantize graph, collect fp32 outputs
            quant_graph.dequantize_graph(expire_device=expire_device)
            for data in dataloader:
                if collate_fn is not None: data = collate_fn(data)
                fp_output = executor.forward(data, [var.name for var in block.ep.outputs])
                fp_output = {var.name: cache_fn(data) for data, var in zip(fp_output, block.ep.outputs)}
                fp_outputs.append(fp_output)