def partial_fit()

in optimum/furiosa/quantization.py [0:0]

• 19 lines of code
• 4 McCabe index (conditional complexity)

    def partial_fit(self, dataset: Dataset, calibration_config: CalibrationConfig, batch_size: int = 1):
        """
        Performs the calibration step and collects the quantization ranges without computing them.

        Args:
            dataset (`Dataset`):
                The dataset to use when performing the calibration step.
            calibration_config (`CalibrationConfig`):
                The configuration containing the parameters related to the calibration step.
            batch_size (`int`, *optional*, defaults to 1):
                The batch size to use when collecting the quantization ranges values.
        """
        self._calibration_config = calibration_config

        # If no calibrator, then create one
        if calibration_config.method is not None:
            LOGGER.info(f"Creating calibrator: {calibration_config.method}({calibration_config})")
            self._load_model_and_optimize()

            self._calibrator = calibration_config.create_calibrator(
                model=self.onnx_model,
            )

        def get_input_datatypes(model):
            input_datatypes = []

            for input in model.graph.input:
                input_type = input.type.tensor_type.elem_type
                input_datatypes.extend([input_type])

            return input_datatypes

        input_datatypes = get_input_datatypes(self.onnx_model)

        LOGGER.info("Collecting tensors statistics...")
        reader = FuriosaAICalibrationDataReader(dataset, input_datatypes, batch_size)
        for data in tqdm.tqdm(reader):
            self._calibrator.collect_data(data)