基于国产异构加速卡的Ultralytics训练实践报告

Ultralytics 计算机视觉框架高度依赖特定版本的PyTorch及其配套的CUDA生态，适配到国产异构加速卡，往往第一步就卡在依赖冲突与算子兼容等问题上。

本次实践，基于超算互联网西南一区的异构加速卡AI （64G 显存），完成 Ultralytics YOLO26 从环境选型、依赖排错、算子源码改造、torchvision 源码编译适配，到模型训练、实时推理、ONNX 模型导出的全流程，逐层对比桩模块过渡、原生包安装、ROCm 源码编译三类适配方案优劣，量化多规格 YOLO 模型在加速卡上的时延、显存性能指标，形成一套 “环境隔离 + 源码编译 + 路径动态挂载” 标准化适配最佳实践。

实践打通了国产异构算力与主流开源目标检测框架的技术壁垒，意味着国产算力生态正从“能用”迈向“好用”完整记录可直接复用的操作流程，为超算互联网平台及行业用户基于国产异构加速卡开展计算机视觉模型训练、推理业务提供详实可落地的工程参考，助力国产 AI 算力生态规模化应用。

1. 配置选型

本次实践基于超算互联网平台 Notebook 开发环境，首先配置最新的pytorch镜像，选择异构加速卡AI（显存64GB），基础镜像PyTorch/2.9.0/py3.11-ubuntu22.04/dtk26.04。PyTorch2.9.0是一个定制版本的PyTorch，与官方PyTorch的ABI（应用二进制接口）不兼容。

2. Ultralytics环境配置

2.1 服务器环境

使用lscpu、hy-smi--showmemuse--showmemavailable、df-h、free-h等命令查看服务器硬件环境。

可见服务器为配备DCU（Deep Computing Unit）的服务器。下面需要在该环境上配置并跑通Ultralytics YOLO项目。 DCU 配套 DTK 软件栈基于 ROCm 架构，与 NVIDIA CUDA 底层接口存在根本性差异。直接使用pip install ultralytics会引入原生CUDA的torch版本。这会覆盖现有的DTK PyTorch，导致环境破坏。因此，需要非常小心地控制pip安装行为，使用--no-deps或手动安装依赖，规避版本覆盖、ABI不兼容问题。

2.2 方案1：安装官方 torchvision（失败）

bash

pip3 install torchvision==0.20.0 --no-deps

结果:安装成功但导入失败，报错operator torchvision::nms not registered，底层算子库与 DTK 架构不匹配，无法正常使用。

2.3 方案2：从PyTorch官方ROCm源安装（失败）

bash

pip3 install torchvision --index-url https://download.pytorch.org/whl/rocm6.3 --no-deps

结果:同样不兼容，ROCm 原生编译包与 DTK 衍生工具链 ABI 不互通，算子库依旧无法加载，推理流程中断。

2.4 方案3：创建torchvision桩模块（成功）

为快速打通基础推理流程，构建虚拟桩模块绕过 TorchVision 强依赖，完整操作流程：

创建伪torchvision元数据

bash

   mkdir -p /usr/local/lib/python3.11/site-packages/torchvision-0.0.0.dist-info
   cat > /usr/local/lib/python3.11/site-packages/torchvision-0.0.0.dist-info/METADATA << 'EOF'
   Metadata-Version: 2.1
   Name: torchvision
   Version: 0.0.0
   Summary: Fake torchvision metadata for Ultralytics
   EOF

创建最小化torchvision桩模块

bash

   # /usr/local/lib/python3.11/site-packages/torchvision/__init__.py
   __version__ = "0.0.0"
   import sys
   
   class _StubModule:
       def __init__(self, name):
           self.__name__ = name
       def __getattr__(self, name):
           return _StubModule(f"{self.__name__}.{name}")
       def __call__(self, *args, **kwargs):
           raise NotImplementedError(f"{self.__name__} is a stub")
   
   _stub_attrs = ["datasets", "io", "models", "ops", "transforms", "utils", "_meta_registrations"]
   for _attr in _stub_attrs:
       if not hasattr(sys.modules[__name__], _attr):
           setattr(sys.modules[__name__], _attr, _StubModule(_attr))

修改Ultralytics NMS 源码文件:/root/ultralytics/ultralytics/utils/nms.py 原代码(第150-156行):

bash

   # Speed strategy: torchvision for val or already loaded (faster), TorchNMS for predict (lower latency)
   if "torchvision" in sys.modules:
       import torchvision  # scope as slow import
       i = torchvision.ops.nms(boxes, scores, iou_thres)
   else:
       i = TorchNMS.nms(boxes, scores, iou_thres)

修改后:

bash

   # Speed strategy: torchvision for val or already loaded (faster), TorchNMS for predict (lower latency)
   # Modified for Hygon DCU: always use TorchNMS.nms to avoid torchvision dependency
   i = TorchNMS.nms(boxes, scores, iou_thres)

原理:

Ultralytics有两种NMS实现:

torchvision.ops.nms (C++实现，更快)
TorchNMS.nms (纯PyTorch实现，较慢但无依赖) 原代码根据"torchvision" in sys.modules选择实现修改后强制使用纯PyTorch实现，避免依赖torchvision

2.5 自行编译torchvision（最优方案）

为完整恢复 DCU 硬件加速 NMS 算子、消除推理性能损耗，本次实践采用本地源码编译 TorchVision作为最终标准化落地方案，依托 DTK 环境编译适配 ROCm 架构的算子库，完整操作流程如下：

2.5.1 首先下载torchvision源码zip文件

2.5.2 编译torchvision开发版

bash

cd /root/ultralytics/vision-main && export ROCM_HOME=/opt/dtk && export FORCE_CUDA=1 && export PYTORCH_VERSION_GE=2.9.0 && export PYTORCH_VERSION_LT=2.10.0 && python3 setup.py build_ext --inplace 2>&1 | tail -200

等待编译成功，生成torchvision/_C.so(3.3M)和torchvision/image.so(431K)

2.5.3 备份stub torchvision

bash

mv /usr/local/lib/python3.11/site-packages/torchvision /usr/local/lib/python3.11/site-packages/torchvision.stub.backup
mv /usr/local/lib/python3.11/site-packages/torchvision-0.0.0.dist-info /usr/local/lib/python3.11/site-packages/torchvision-0.0.0.dist-info.backup

等待备份完成

2.5.4 安装编译后的torchvision（创建.pth文件）

bash

echo "/root/ultralytics/vision-main" > /usr/local/lib/python3.11/site-packages/vision-main.pth

.pth 文件创建成功

2.5.5 创建torchvision元数据（解决PackageNotFoundError）

bash

mkdir -p /usr/local/lib/python3.11/site-packages/torchvision-0.28.0a0.dist-info
cat > /usr/local/lib/python3.11/site-packages/torchvision-0.28.0a0.dist-info/METADATA << 'EOF'
Metadata-Version: 2.1
Name: torchvision
Version: 0.28.0a0+96d1d17
Summary: Image and video datasets and models for torch deep learning
EOF

元数据创建成功

2.5.6 测试torchvision导入

bash

python3 -c "import torchvision; print('version:', torchvision.__version__); print('location:', torchvision.__file__)"

输出: version:0.28.0a0+96d1d17 location:/root/ultralytics/vision-main/torchvision/init.py torchvision导入成功

2.5.7 测试torchvision.ops.nms

bash

python3 -c "
import torch
import torchvision
boxes = torch.tensor([[100, 100, 200, 200], [150, 150, 250, 250]], dtype=torch.float32).cuda()
scores = torch.tensor([0.9, 0.8], dtype=torch.float32).cuda()
indices = torchvision.ops.nms(boxes, scores, 0.5)
print('NMS indices:', indices)
print('SUCCESS!')
"

输出: NMS indices:tensor([0,1],device='cuda:0') SUCCESS! 结果:NMS在DCU上工作正常

2.5.8 恢复ultralytics/utils/nms.py（改回使用torchvision.ops.nms）

bash

# 修改文件: /root/ultralytics/ultralytics/utils/nms.py
# 将第 150-152 行从:
#   # Modified for Hygon DCU: always use TorchNMS.nms to avoid torchvision dependency
#   i = TorchNMS.nms(boxes, scores, iou_thres)
# 改为:
#   if "torchvision" in sys.modules:
#       import torchvision
#       i = torchvision.ops.nms(boxes, scores, iou_thres)
#   else:
#       i = TorchNMS.nms(boxes, scores, iou_thres)

nms.py恢复成功

2.5.9 测试Ultralytics推理（使用编译后的torchvision）

bash

cd /root/ultralytics && python3 -c "
from ultralytics import YOLO
model = YOLO('weights/yolo26n.pt')
results = model('ultralytics/assets/zidane.jpg', save=True, project='runs/detect', name='zidane_torchvision_test')
print('Results:', len(results), 'images')
for r in results:
    print('Detections:', len(r.boxes), 'boxes')
"

2.6 清理损坏的包元数据

目标:清理site-packages中损坏的包元数据，使pip能正常工作执行内容:

bash

rm -rf /usr/local/lib/python3.11/site-packages/torchvision
rm -rf /usr/local/lib/python3.11/site-packages/torchvision-0.28.0a0-py3.11.egg-info
rm -rf /usr/local/lib/python3.11/site-packages/scipy
rm -rf /usr/local/lib/python3.11/site-packages/scipy-1.17.1.dist-info
rm -rf /usr/local/lib/python3.11/site-packages/polars
rm -rf /usr/local/lib/python3.11/site-packages/polars-1.40.1.dist-info

2.7 安装Ultralytics依赖

目标:安装Ultralytics项目所需的核心依赖包执行内容:

干运行检查依赖

bash

   pip3 install scipy polars ultralytics-thop --dry-run

结果:只会安装scipy-1.17.1,polars-1.40.1,ultralytics-thop，无torch或CUDA包

安装依赖

bash

   pip3 install scipy polars ultralytics-thop

验证安装

bash

   import scipy; print('scipy:', scipy.__version__)  # 1.17.1
   import polars; print('polars:', polars.__version__)  # 1.40.1

已安装的依赖包: numpy 1.26.4、opencv-python 4.13.0、pillow 12.2.0、pyyaml 6.0.3、matplotlib 3.10.8、scipy 1.17.1、psutil 7.2.2、polars 1.40.1、ultralytics-thop 2.0.19

2.8 配置Ultralytics项目

目标:使Ultralytics项目可被Python导入，无需pip install -e . 问题分析: pip install -e .会升级torch到2.12.0并安装大量CUDA包这会破坏DCU环境。解决方案:使用.pth文件执行内容:

bash

echo "/root/ultralytics" > /usr/local/lib/python3.11/site-packages/ultralytics.pth

原理: Python在启动时会读取site-packages目录下的所有.pth文件 .pth文件中的每一行都会被添加到sys.path 这样/root/ultralytics就被永久添加到Python路径中验证:

bash

from ultralytics import YOLO  # 成功导入

执行测试：

2.9 测试推理功能

目标: 验证Ultralytics项目能在DCU上正常推理执行内容:

准备测试图片使用Ultralytics自带的测试图片: /root/ultralytics/ultralytics/assets/zidane.jpg
运行推理

bash

   from ultralytics import YOLO
   model = YOLO('yolo26l.pt')
   results = model('ultralytics/assets/zidane.jpg')

查看结果

bash

   print('Results:', len(results), 'images processed')
   [print(r.boxes) for r in results]

输出: image 1/1 /root/ultralytics/ultralytics/assets/bus.jpg: 640x480 4 persons, 1 bus, 48.8ms Speed: 3.4ms preprocess, 48.8ms inference, 1.0ms postprocess per image at shape (1, 3, 640, 480)

保存推理结果

bash

   results = model('ultralytics/assets/zidane.jpg', save=True)
   print('Results saved to:', results[0].save_dir)

结果位置:/root/ultralytics/runs/detect/predict/zidane.jpg 同样加载n，s，m，l，x模型进行推理，对比结果：

2.10 执行训练功能测试

手动下载数据集文件： https://github.com/ultralytics/assets/releases/download/v0.0.0/coco8.zip 解压到ultralytics目录下，启动训练：

bash

yolo train model=weights/yolo26n.pt data=ultralytics/cfg/datasets/coco8.yaml
epochs=10 imgsz=640 batch=4 device=0

训练过程正常，保存模型文件正常。

3. YOLO26n ONNX推理测试最终报告

3.1 ONNX模型导出

导出命令:

bash

cd /root/ultralytics
yolo export model=weights/yolo26n.pt format=onnx opset=17

导出结果: 模型已导出:/root/ultralytics/weights/yolo26n.onnx 文件大小:9.5 MB 导出时间:19.3s ONNX opset:17

模型输入/输出: 输入名称:images 输入形状:[1,3,640,640] 输出名称:output0 输出形状:(1,300,6)←[batch,num_detections,x1y1x2y2confcls]

3.2 使用ONNX CPU推理

bash

import onnxruntime as ortsession = ort.InferenceSession('weights/yolo26n.onnx', providers=['CPUExecutionProvider'])

3.3 测试成果

ONNX模型导出成功模型:YOLO26n→ONNX(opset=17) 大小:9.5MB 路径:/root/ultralytics/weights/yolo26n.onnx
ONNX推理测试成功检测对象:3个(2人+1领带) 置信度:Person1=0.918,Person2=0.903,Tie=0.530 与PyTorch对比:置信度差异<1%，结果完全一致！
性能对比分析指标 PyTorch(DCU) ONNX(CPU) 差异推理时间 47.96ms 707.43ms 14.7x慢检测精度 3个对象 3个对象一致置信度 0.915/0.910/0.527 0.918/0.903/0.530 <1%差异效果对比： pytorch onnx（cpu）

4. 总结

本次实践，我们基于超算互联网异构加速卡，通过系统化排查与工程化手段攻克了精细化控制pip依赖安装、利用.pth文件动态挂载项目路径、先期构建torchvision桩模块过渡等问题，并最终完成torchvision源码级编译与深度适配，成功实现了YOLO模型推理、训练与ONNX导出全流程的稳定运行。尽管DCU与主流CUDA生态在底层接口上存在差异，但通过“桩模块过渡+源码编译+环境精细化管控”等系统化方法，能够完全打通技术壁垒，使Ultralytics等高度依赖CUDA特性（如torchvision、自定义NMS等）的主流开源项目在DCU上实现无缝兼容与高效运行。下一步，我们将持续聚焦推理性能调优与训练流程标准化，并在更多主流视觉与多模态模型上开展迁移验证，进一步夯实DCU在复杂AI工作流中的稳定性与可移植性，为国产算力生态的规模化应用提供实践参考。

提交作业

使用指导

CAE仿真

科学计算

最佳实践

Notebook快速使用

Notebook功能介绍

应用类

开发类

购买商品

最佳实践案例

基于国产异构加速卡的Ultralytics训练实践报告 ​

1. 配置选型 ​

2. Ultralytics环境配置 ​

2.1 服务器环境 ​

2.2 方案1：安装官方 torchvision（失败） ​

2.3 方案2：从PyTorch官方ROCm源安装（失败） ​

2.4 方案3：创建torchvision桩模块（成功） ​

2.5 自行编译torchvision（最优方案） ​

2.6 清理损坏的包元数据 ​

2.7 安装Ultralytics依赖 ​

2.8 配置Ultralytics项目 ​

2.9 测试推理功能 ​

2.10 执行训练功能测试 ​

3. YOLO26n ONNX推理测试最终报告 ​

3.1 ONNX模型导出 ​

3.2 使用ONNX CPU推理 ​

3.3 测试成果 ​

4. 总结 ​

基于国产异构加速卡的Ultralytics训练实践报告

1. 配置选型

2. Ultralytics环境配置

2.1 服务器环境

2.2 方案1：安装官方 torchvision（失败）

2.3 方案2：从PyTorch官方ROCm源安装（失败）

2.4 方案3：创建torchvision桩模块（成功）

2.5 自行编译torchvision（最优方案）

2.6 清理损坏的包元数据

2.7 安装Ultralytics依赖

2.8 配置Ultralytics项目

2.9 测试推理功能

2.10 执行训练功能测试

3. YOLO26n ONNX推理测试最终报告

3.1 ONNX模型导出

3.2 使用ONNX CPU推理

3.3 测试成果

4. 总结