openvino 量化

要实现Openvino训练后量化，需要两步：

1. 将Onnnx文件转成Openvino格式。

2.将上述文件利用POT进行量化，在量化过程中，我们需要提供量化验证数据集和验证指标以保证量化后的效果。

第一步：

MODEL_DIR.mkdir(exist_ok=True)
!mo --input_model $ONNX_PATH --output_dir $MODEL_DIR --compress_to_fp16 --scale_values [1] --mean_values [0]

第二步：

class BinaryF1(Metric):
    """
    Metric to compute F1/Dice score for binary segmentation. F1 is computed as
    (2 * precision * recall) / (precision + recall) where precision is computed as
    the ratio of pixels that were correctly predicted as true and all actual true pixels,
    and recall as the ratio of pixels that were correctly predicted as true and all
    predicted true pixels.

    See https://en.wikipedia.org/wiki/F-score
    """

    # Required methods
    def __init__(self):
        super().__init__()
        self._name = "F1"
        self.y_true = 0
        self.y_pred = 0
        self.correct_true = 0

    @property
    def value(self):
        """
        Returns metric value for the last model output.
        Possible format: {metric_name: [metric_values_per_image]}
        """
        return {self._name: [0, 0]}

    @property
    def avg_value(self):
        """
        Returns average metric value for all model outputs as {metric_name: metric_value}
        """
        recall = self.correct_true / self.y_pred
        precision = self.correct_true / self.y_true

        f1 = (2 * precision * recall) / (precision + recall)
        return {self._name: f1}

    def update(self, output, target):
        """
        Update the statistics for metric computation

        :param output: model output
        :param target: annotations for model output
        """
        label = target[0].astype(np.byte)
        prediction = output[0].astype(np.byte)

        self.y_true += np.sum(label)
        self.y_pred += np.sum(prediction)

        correct_true = np.sum(
            (label == prediction).astype(np.byte) * (label == 1).astype(np.byte)
        ).astype(np.float32)

        self.correct_true += correct_true

    def reset(self):
        """
        Resets metric to initial values
        """
        self.y_true = 0
        self.y_pred = 0
        self.correct_true = 0

    def get_attributes(self):
        """
        Returns a dictionary of metric attributes {metric_name: {attribute_name: value}}.
        Required attributes: 'direction': 'higher-better' or 'higher-worse'
                             'type': metric type
        """
        return {self._name: {"direction": "higher-better", "type": "F1"}}

def rotate_and_flip(image):
    """Rotate `image` by 90 degrees and flip horizontally"""
    return cv2.flip(cv2.rotate(image, rotateCode=cv2.ROTATE_90_CLOCKWISE), flipCode=1)


class KitsDataLoader(DataLoader):
    def __init__(self, basedir: str):
        """
        DataLoader class for prepared KiTS19 data, for binary segmentation (background/kidney)
        Source data should exist in basedir, in subdirectories case_00000 until case_00210,
        with each subdirectory containing directories imaging_frames, with jpg images, and
        segmentation_frames with segmentation masks as png files.
        See https://github.com/openvinotoolkit/openvino_notebooks/blob/main/notebooks/110-ct-segmentation-quantize/data-preparation-ct-scan.ipynb

        For demonstration purposes, in this implementation we use images from the validation subset.

        :param basedir: Directory that contains the prepared CT scans
        """
        allmasks = sorted(BASEDIR.glob("case_*/segmentation_frames/*png"))

        if len(allmasks) == 0:
            raise ValueError(f"basedir: '{basedir}' does not contain data")
        val_indices = []  # fmt: skip
        val_cases = []
        masks = [mask for mask in allmasks]

        self.basedir = basedir
        self.dataset = masks
        print(
            f"Created dataset with {len(self.dataset)} items. "
            f"Base directory for data: {basedir}"
        )

    def __getitem__(self, index):
        """
        Get an item from the dataset at the specified index.

        :return: (annotation, input_image, metadata) where annotation is (index, segmentation_mask)
                 and metadata a dictionary with case and slice number
        """
        mask_path = self.dataset[index]
        image_path = str(mask_path.with_suffix(".jpg")).replace(
            "segmentation_frames", "imaging_frames"
        )

        # Load images with MONAI's LoadImage to match data loading in training notebook
        mask = LoadImage(image_only=True, dtype=np.uint8)(str(mask_path)).numpy()
        img = LoadImage(image_only=True, dtype=np.float32)(str(image_path)).numpy()
        annotation = (index, mask[56:56+400,56:56+400])
        input_image=np.expand_dims(img[56:56+400,56:56+400],axis=0)/255
#         input_image=np.expand_dims(input_image,axis=0)
        return (
            annotation,
            input_image,
            {"case": Path(mask_path).parents[1].name, "slice": Path(mask_path).stem},
        )

    def __len__(self):
        return len(self.dataset)

data_loader = KitsDataLoader(BASEDIR)

# Step 2: Load the model.
ir_model = load_model(model_config=model_config)

# Step 3: Initialize the metric.
metric = BinaryF1()

# Step 4: Initialize the engine for metric calculation and statistics collection.
engine = IEEngine(config=engine_config, data_loader=data_loader, metric=metric)

# Step 5: Create a pipeline of compression algorithms.
# The `quantization_algorithm` parameter is defined in the Settings.
pipeline = create_pipeline(algorithms, engine)

# Step 6: Execute the pipeline to quantize the model.
algorithm_name = pipeline.algo_seq[0].name
preset = pipeline._algo_seq[0].config["preset"]
print(f"Executing POT pipeline on {model_config['model']} with {algorithm_name}, {preset} preset")
start_time = time.perf_counter()
compressed_model = pipeline.run(ir_model)
end_time = time.perf_counter()
print(f"Quantization finished in {end_time - start_time:.2f} seconds")

# Step 7 (Optional): Compress model weights to quantized precision
#                    in order to reduce the size of the final .bin file.
compress_model_weights(compressed_model)

# Step 8: Save the compressed model to the desired path.
# Set `save_path` to the directory where the model should be saved.
compressed_model_paths = save_model(
    model=compressed_model,
    save_path="optimized_model",
    model_name=f"{ir_model.name}_{preset}_{algorithm_name}",
)
compressed_model_path = compressed_model_paths[0]["model"]
print("The quantized model is stored at", compressed_model_path)