使用 NNCF 量化模型（Python篇）

简介

本文主要介绍 OpenVINO 量化中常用的训练后量化的基本量化流程，训练后量化就是对训练好的模型进行量化的一种方式，这种方式相对来说比较常用，也更简单，只需要准备和模型输入相同的小规模的校准数据集就可以。其他量化过程请参考官方介绍。这里以图像数据为例，量化图像编码器模型。

一、优化工具安装

1. 在 OpenVINO 中，默认的优化工具是 NNCF，但 NNCF 默认不包含在 OpenVINO 工具箱中，所以需要在我们的环境中通过pip进行安装。

   pip install nncf
   

2. NNCF（神经网络压缩框架）支持OpenVINO IR、PyTorch、ONNX 和 TensorFlow 格式的模型，主要提供权重压缩、训练后量化和训练时优化。更多细节可以参考官方教程和具体实例，这里不过多赘述。

二、准备校准数据集

1. 在使用 NNCF 做训练后量化（Post-Training Quantization） 时，通常需要提供一个数据集来统计模型在“真实数据”下的激活分布。这里以COCO数据集为例，构建数据集类，指定图像的加载逻辑，因为是用于量化时的校准输入，所以不需要标签，之后被用于 DataLoader 中。

   import torch
   import torch.utils.data as dataclass COCOLoader(data.Dataset):def __init__(self, images_path):self.images = list(Path(images_path).iterdir())def __getitem__(self, index):image_path = self.images[index]image = cv2.imread(str(image_path))image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)return imagedef __len__(self):return len(self.images)
   

2. 创建一个数据集实例并初始化刚才定义的数据集类 COCOLoader。

   coco_dataset = COCOLoader("./coco128/images/train2017")
   

3. 创建一个数据加载器，按 batch_size 自动加载数据。在量化过程中 calibration_loader 负责把数据一张张送入模型，让 NNCF 收集模型中各层的激活范围，生成量化感知的模型。通常设置每次只加载 1 张图像。

   calibration_loader = data.DataLoader(coco_dataset，batch_size=1)  # 默认 batch_size=1
   

4. 上面我们已经定义了从哪里拿以及怎么拿数据，这其实和训练模型的时候一样。同样的，这些数据在进入模型之前通常会做预处理，所以在执行真正的量化之前，需要定义一个转换函数，转换函数是一种从数据集中获取样本并返回可传递给模型进行推理的数据的函数。需要结合模型的实际情况定义预处理过程，这里的例子仅供参考。

   def transform_fn(image_data):if isinstance(image_data, (torch.Tensor, np.ndarray)) and image_data.ndim == 4:image_data = image_data[0].cpu().numpy()  # 去除 batch 维度，变成 3D 图像_, _, input_image = processor.set_image(image_data)  # processor是自己定义的，这里仅作参考input_image = input_image.cpu().numpy()return input_image
   

5. 这样，我们就定义好了如何拿数据，以及拿到数据之后要做怎么的预处理，之后就可以调用NNCF 具体的api加载数据并进行量化。

三、量化模型和保存

1. 首先，创建一个用于量化校准（Calibration）的数据集包装器，它会将 calibration_loader 中加载出来的每个数据样本，经过 transform_fn 处理后再输入到模型中，从而让 NNCF 获取合适的输入格式和范围，用于计算量化参数（如激活值范围）。这是量化流程中一个关键步骤，它负责包装你加载的数据和预处理逻辑，确保数据可以正确地送入模型进行量化校准。

   import nncfcalibration_dataset = nncf.Dataset(calibration_loader, transform_fn)
   

2. 之后就可以调用 quantize 方法对一个预训练模型进行训练后量化，需要指定待量化的模型（model）和校准数据集（calibration_dataset），并且可以通过model_type参数告诉 NNCF 模型的类型。NNCF 会根据模型结构应用最合适的量化策略，比如指定CNN：会用标准卷积量化策略，TRANSFORMER：使用适合注意力机制的量化方法（如保留 LayerNorm、限制权重量化范围等）。为了加速量化，可以通过subset_size参数选择样本的子集，而不需要遍历整个数据集，通常取几十到几百张图像就够用了。

   import Path
   import openvino as ovcore = ov.Core()
   model_path=Path("./model.xml")model = core.read_model(model_path)
   quantized_model = nncf.quantize(model,calibration_dataset,model_type=nncf.parameters.ModelType.TRANSFORMER,subset_size=128,
   )
   

3. 最后，将量化后的模型保存到指定路径下。

   quantized_path = Path("./model_int8.xml")
   ov.save_model(quantized_model, quantized_path)
   

小结

这是一个很简单的例子，实际使用的时候可能会因为数据集以及预处理方式等问题，使情况变得复杂一点点，所以更重要的还是看官方的样例。以上分享均来自个人实践和官网介绍，如有其他问题欢迎在评论区讨论，谢谢！！！