1. 项目概述
在计算机视觉领域,目标检测一直是核心任务之一。YOLO系列模型因其出色的实时性和准确性而广受欢迎。最近,我在Atlas 300I推理卡上成功部署了YOLOv26模型,并通过ATC工具将PyTorch模型转换为昇腾专用的OM格式。整个过程涉及环境搭建、模型转换、推理验证和精度评估等多个环节,最终实现了OM模型与原始PT模型在mAP指标上仅千分之一的误差。
这个项目对于需要在昇腾AI处理器上部署目标检测模型的研究人员和开发者具有重要参考价值。它不仅验证了昇腾310芯片对YOLO系列模型的兼容性,还提供了一套完整的模型转换和验证流程。
需要模型API调用? 免费领10W Token,多模型网关一键接入 Claude、DeepSeek 等主流模型。
2. 环境准备
2.1 硬件配置
本次项目使用的硬件平台是华为Atlas 300I推理卡,搭载昇腾310 AI处理器。这款处理器专为边缘推理场景设计,具有以下特点:
- 算力:16TOPS INT8
- 功耗:8W
- 支持FP16/INT8混合精度计算
- 内置DVPP硬件加速模块
2.2 软件环境搭建
完整的软件栈包括:
code复制操作系统:Ubuntu 20.04 LTS
驱动版本:A300-3010-npu-driver_24.1.1.3
固件版本:A300-3010-npu-firmware_7.5.0.9.220
CANN工具包:Ascend-cann-toolkit_8.0.1
Python环境:3.8.10
PyTorch:2.9.1
Ultralytics:8.4.3
安装步骤:
- 安装驱动和固件:
bash复制chmod +x A300-3010-npu-driver_24.1.1.3_linux-aarch64.run
./A300-3010-npu-driver_24.1.1.3_linux-aarch64.run --full
chmod +x A300-3010-npu-firmware_7.5.0.9.220.run
./A300-3010-npu-firmware_7.5.0.9.220.run --full
- 安装CANN工具包:
bash复制chmod +x Ascend-cann-toolkit_8.0.1_linux-aarch64.run
./Ascend-cann-toolkit_8.0.1_linux-aarch64.run --install
注意:安装过程中需要确保系统已安装必要的依赖库,如libssl、libsqlite等。建议使用root权限安装,并按照官方文档配置环境变量。
3. 模型转换流程
3.1 PyTorch转ONNX
YOLOv26模型首先需要从PyTorch格式(.pt)转换为ONNX格式。转换脚本关键点:
python复制import argparse
from ultralytics import YOLO
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--pt', default="./models/yolo26s.pt", help='pt file')
args = parser.parse_args()
model = YOLO(args.pt)
onnx_model = model.export(format="onnx", dynamic=True, simplify=True, opset=11)
if __name__ == '__main__':
main()
转换命令:
bash复制python3 pth2onnx.py --pt ./models/yolo26s.pt
转换参数说明:
dynamic=True:允许动态batch sizesimplify=True:启用ONNX简化opset=11:使用ONNX opset 11版本
3.2 ONNX转OM
使用ATC工具将ONNX模型转换为昇腾OM格式:
bash复制atc --framework=5 \
--model=./models/yolo26s.onnx \
--input_format=NCHW \
--input_shape='images:1,3,640,640' \
--output_type=FP32 \
--output=./models/yolo26s_bs1 \
--log=info \
--soc_version=Ascend310
关键参数解析:
--framework=5:指定输入模型为ONNX格式--input_shape:设置模型输入尺寸为1×3×640×640--output_type=FP32:输出数据类型为FP32--soc_version=Ascend310:指定目标芯片型号
实际项目中,建议先使用小batch size(如1)进行转换验证,成功后再尝试更大batch size。转换过程中遇到问题可以通过
--log=debug获取更详细日志。
4. 模型推理实现
4.1 推理代码解析
推理脚本实现了PT模型和OM模型的对比验证,核心逻辑如下:
python复制def pt_detect(input_args):
model = YOLO(input_args.pt)
output = model(source=input_args.data, save=True, save_txt=True)
# 保存结果到文本文件
def om_detect(input_args):
om_model = InferSession(int(input_args.device_id), input_args.om)
dp = DetectionPredictor(overrides=om_conf.overrides)
# 预处理、推理、后处理流程
preds = om_model.infer([im])
# 保存结果到文本文件
关键点说明:
- PT模型直接使用Ultralytics的YOLO接口
- OM模型通过ais_bench的InferSession进行推理
- 两种方式的结果都保存为相同格式的文本文件,便于后续对比
4.2 推理结果验证
执行推理命令:
bash复制python3 inference.py --data ./val2017 --pt ./models/yolo26s.pt --om ./models/yolo26s_bs1.om
验证结果显示PT和OM模型的检测框位置、类别和置信度高度一致,视觉上几乎无法区分:
code复制PT模型检测结果: [tensor([[ 68.12500, 344.50000, 116.37500, 383.25000, 0.87695, 0.00000],
[331.62500, 344.50000, 380.12500, 383.25000, 0.86328, 0.00000]])]
OM模型检测结果: [tensor([[ 68.12305, 344.49805, 116.37695, 383.25195, 0.87690, 0.00000],
[331.62695, 344.49805, 380.12305, 383.25195, 0.86320, 0.00000]])]
5. 精度评估方法
5.1 mAP计算原理
mAP(mean Average Precision)是目标检测领域最常用的评估指标。本项目采用COCO官方评估标准,主要计算以下指标:
- AP@[0.50:0.95]:IoU阈值从0.5到0.95,步长0.05的平均AP
- AP@0.50:IoU阈值为0.5时的AP
- AP@0.75:IoU阈值为0.75时的AP
- APsmall/medium/large:不同尺度目标上的AP
5.2 评估代码实现
使用pycocotools计算mAP的核心代码:
python复制def eval_compute(cocoDt_json,cocoGt_file):
cocoGt = COCO(cocoGt_file)
cocoDt = cocoGt.loadRes(cocoDt_json)
cocoEval = COCOeval(cocoGt, cocoDt, "bbox")
cocoEval.evaluate()
cocoEval.accumulate()
cocoEval.summarize()
执行评估:
bash复制# PT模型评估
python3 eval_map.py --xml_dir pt-result/
# OM模型评估
python3 eval_map.py --xml_dir om-result/
5.3 评估结果对比
YOLOv26各版本模型的精度对比:
| 模型 | PT格式mAP | OM格式mAP | 误差 |
|---|---|---|---|
| YOLO26n | 0.330 | 0.330 | 0.000 |
| YOLO26s | 0.406 | 0.405 | 0.001 |
| YOLO26m | 0.455 | 0.454 | 0.001 |
| YOLO26l | 0.469 | 0.469 | 0.000 |
| YOLO26x | 0.499 | 0.498 | 0.001 |
结果显示所有模型的转换误差都控制在千分之一以内,证明了ATC工具转换的可靠性。
6. 关键问题与解决方案
6.1 输入输出对齐
问题现象:OM模型输出与PT模型存在微小差异
解决方案:
- 确保ATC转换时的输入输出数据类型一致(FP32)
- 在推理代码中对OM输出做相同的后处理
- 对输出结果进行四舍五入处理,消除浮点误差
6.2 标签映射问题
问题现象:COCO数据集的类别ID与YOLO不一致
解决方案:建立映射表转换类别ID
python复制CATEGORY_ID = {
0:1, 1:2, 2:3, 3:4, 4:5, 5:6, 6:7, 7:8, 8:9, 9:10,
# ...完整映射表
79:90
}
6.3 预处理一致性
问题现象:letterbox处理方式不同导致输入差异
解决方案:统一预处理逻辑
python复制def patch_pre_transform(self, im):
same_shapes = len({x.shape for x in im}) == 1
self.model.pt = False
letterbox = LetterBox(self.imgsz, auto=same_shapes, stride=self.model.stride)
return [letterbox(image=x) for x in im]
7. 性能优化建议
-
批量推理优化:
- 尝试更大的batch size(如4/8/16)
- 使用
--input_shape='images:4,3,640,640'转换模型 - 合理设置
--output_type=FP16减少内存占用
-
后处理加速:
- 将NMS等后处理操作移到模型内
- 使用C++插件实现自定义算子
-
量化部署:
- 尝试INT8量化减少模型大小
- 使用ATC的
--quantize参数进行量化
-
多卡并行:
- 对于大batch size场景,可以使用多卡并行推理
- 通过
device_id参数指定不同卡
在实际部署中发现,对于640×640的输入尺寸,Atlas 300I单卡可以稳定运行batch size=4的FP16模型,帧率能达到45FPS,完全满足实时检测需求。
