YOLOv8 分类模型用 onnx 推理的前处理

Mofify Date: 2024-07-07

场景描述

跟随 YOLOv8 推出的是 ultralytics 训练和推理框架，可以一行代码实现分类、检测、分割等任务的训练，以及推理和模型导出。

发现在使用默认设置训练分类模型，并使用 ultralytics 库将模型 pt 权重一键导出为 onnx 后，如果不使用 ultralytics 加载 onnx 权重，而直接在其他没有 ultralytics 的环境和设备上用 onnxruntime 直接加载权重并推理，会出现推理结果不一致甚至置信度很低的情况。这是缺少 pt 权重文件所能表达的前处理信息，没有实现前处理而导致的。需要补充 YOLOv8 自带的前处理。

---

解决方案

可以通过对 pt 文件的前处理部分解析，或者对 ultralytics 库中 YOLO 分类任务的预测部分代码的源码解析（ultralytics/models/yolo/classify/predict.py），找到前处理的代码，补充到 onnxruntime 推理代码之前即可：

import torchvision.transforms as T
from torchvision.transforms import InterpolationMode

def get_classify_transforms(size: int=224, mean: Tuple[float, float, float]=(0.0, 0.0, 0.0),
                            std: Tuple[float, float, float]=(1.0, 1.0, 1.0),
                            interpolation: InterpolationMode=InterpolationMode.BILINEAR,
                            crop_fraction: float = 1.0) -> T.Compose:
    """
    复现 YOLO 默认的分类预处理
    """
    if isinstance(size, (tuple, list)):
        assert len(size) == 2
        scale_size = tuple(math.floor(x / crop_fraction) for x