快速开始

本章以一个 YOLO 模型为示例，讲解 TVM-MUSA 的编译与推理流程。示例拆成两个独立阶段：

• build.py：使用 Ultralytics、Torch-MUSA 和 TorchVision 环境加载、融合并导出 YOLOv8s，再通过 TVM-MUSA 编译和优化，最终保存 .so 共享库及描述文件。
• inference.py：只加载 .so 和描述文件，在 MUSA 上执行推理、图片后处理、画框和性能测试，不依赖 PyTorch 或 Ultralytics。

这种拆分把构建依赖和部署依赖隔离开。构建阶段可以在具备 Torch-MUSA、TorchVision、Ultralytics 和 TVM-MUSA 的环境中完成；部署阶段需要 TVM-MUSA、NumPy 和 Pillow。

阶段	脚本	主要依赖	产物
编译优化	build.py	TVM-MUSA、Torch-MUSA、TorchVision、Ultralytics、NumPy	yolov8s_musa.so、yolov8s_musa.json
运行时推理	inference.py	TVM-MUSA、NumPy、Pillow	标注图片、检测结果、性能数据

开始前请先完成安装与验证。TVM MUSA Wheel 包以及示例脚本可从从产品下载页面下载

准备示例环境

构建示例需要使用 Moore Threads 发布的 PyTorch 软件栈。请从 M1000 软件发布包中获取与本地 MUSA SDK 匹配的 aarch64 wheel，并将 torch、torch_musa 和 torchvision 一起安装到安装与验证创建的 venv：

python -m pip install \
  /path/to/torch-<version>-cp310-cp310-linux_aarch64.whl \
  /path/to/torch_musa-<version>-cp310-cp310-linux_aarch64.whl \
  /path/to/torchvision-<version>-cp310-cp310-linux_aarch64.whl

本文验证组合为 Torch 2.9.1.post1、Torch-MUSA 2.9.1.post1 和 TorchVision 0.22.1

验证 Torch-MUSA 环境：

python -c "import torch, torch_musa; print(torch.musa.is_available())"

输出为 True 表示环境可用。安装示例的其他依赖：

python -m pip install numpy pillow==12.3.0 ultralytics==8.4.192

Step1 编译并生成 .so 文件

进入发布包的 examples/ 目录，执行构建脚本：

cd tvm-musa-m1000/examples
python build.py

脚本的位置参数是 Ultralytics 模型名，默认 yolov8s.pt，并生成：

• artifacts/<model>_musa.so：TVM-MUSA 编译产物，默认模型对应 artifacts/yolov8s_musa.so；
• artifacts/<model>_musa.json：输入 dtype、图像尺寸和类别名等运行时元数据。

模型文件不存在时，Ultralytics 会自动下载到当前工作目录。也可以直接指定其他 Yolo 模型名：

python build.py yolov8n.pt

离线设备需要提前准备对应模型文件。构建阶段固定 batch=1、640x640 静态输入，模型和输入均为 FP16，输出也是 FP16。默认模型实测输出示例：

PASS: /path/to/tvm-musa-m1000/examples/artifacts/yolov8s_musa.so

构建过程直接使用内存中的 PyTorch ExportedProgram 作为 TVM 前端输入，不会生成 ONNX 或其他中间模型文件。

.so 是面向 M1000、MUSA SDK 和当前 TVM-MUSA 版本生成的平台产物。更换设备、MUSA SDK 或 TVM wheel 后，应重新执行 build.py，不要直接复用其他环境的 .so。

Step2 加载 .so 执行推理

保持在同一目录，直接运行运行时脚本：

python inference.py

inference.py 默认加载以下文件：

• artifacts/yolov8s_musa.so：由 build.py 生成；
• artifacts/yolov8s_musa.json：与 .so 同名的描述文件；
• assets/bus.jpg：用于实际推理和计时的图片。

如果构建时使用了其他模型名，请通过 --model artifacts/<model>_musa.so 指定对应产物。

标注结果自动保存到输入图片所在目录，文件名为 <原文件名>_marked<后缀>，例如 assets/bus_marked.jpg。不需要额外传入输出路径。

在 M1000、MUSA SDK 5.2.0 上使用默认参数时，实测输出示例：

PASS: 5 detections, 10.38 ms, 96.4 images/s
result: /path/to/tvm-musa-m1000/examples/assets/bus_marked.jpg

延迟和吞吐会随设备负载和模型版本波动。上述数值用于说明输出格式和计时口径，不作为固定性能承诺。

以下条件同时满足，表示两阶段链路已经跑通：

• 构建阶段生成了 .so 和同名 .json；
• 标注图片保存在输入图片旁边，检测结果合理；
• 两个脚本均以状态码 0 退出，并分别打印构建或推理的 PASS。

Step3 脚本做了什么

1. 加载、融合并导出 YOLOv8s

build.py 把命令行传入的模型名直接交给 Ultralytics；模型不存在时会自动下载。然后执行常见的 Conv-BatchNorm 融合：

from ultralytics import YOLO

yolo = YOLO(model_name)
fused = yolo.model.fuse(verbose=False)
model = (fused if fused is not None else yolo.model).eval().half()

随后把 YOLOv8 Detect 头切换到静态导出模式，并使用固定 shape 的输入执行一次前向：

import torch
from ultralytics.nn.modules import Detect

for module in model.modules():
    if isinstance(module, Detect):
        module.dynamic = False
        module.export = True
        module.shape = None

example_input = torch.zeros(1, 3, 640, 640, dtype=torch.float16)
with torch.no_grad():
    model(example_input)
exported_program = torch.export.export(
    model, args=(example_input,), strict=True
)

这里固定 batch=1、640x640 和静态 shape。Detect 头默认会根据输入重新生成 anchor；提前设置静态导出模式并执行一次固定输入前向，可以把 anchor 和 stride 固定下来，避免 torch.export 在 make_anchors 中遇到未绑定的动态符号。

2. 导入 TVM 并编译

import tvm
from tvm import relax
from tvm.relax.frontend.torch import from_exported_program

target = tvm.target.Target("musa")
mod = from_exported_program(exported_program, keep_params_as_input=False)

pipe = relax.get_pipeline("musa_mp22_cv", target=target)
executable = tvm.compile(mod, target=target, relax_pipeline=pipe)

executable.export_library("artifacts/yolov8s_musa.so")

1. from_exported_program 将 PyTorch ExportedProgram 转换为 Relax IRModule。导入时默认先执行 PyTorch 算子分解，再将 FX 图节点映射为 Relax 算子；keep_params_as_input=False 会把权重绑定为模块常量，运行时只传输入张量。模型和输入已在导出前转为 FP16，因此生成的 Relax 图保持 FP16。

2. relax.get_pipeline("musa_mp22_cv", target) 获取面向 M1000 / MP22 静态 shape CV 模型的编译 pipeline。tvm.compile 在 MUSA target 下执行该 pipeline，完成图融合、NHWC 布局转换、muDNN 子图划分、算子展开、TIR 调度和内存规划，并生成 host 代码与 MUSA kernel，返回可由 VM 加载的 Executable。

3. executable.export_library() 将 host 代码、MUSA kernel 和依赖的 runtime 模块打包为单个 .so，后续可通过 tvm.runtime.load_module 加载。构建脚本还会生成同名 .json，保存输入 dtype、图像尺寸和类别名，供推理脚本读取。

3. 加载共享库并执行推理

运行时脚本只加载编译产物，不依赖 PyTorch 或 Ultralytics：

import tvm
from tvm import relax

dev = tvm.musa(0)
library = tvm.runtime.load_module("artifacts/yolov8s_musa.so")
vm = relax.VirtualMachine(library, dev)

input_tensor = tvm.runtime.tensor(input_np, dev)
output = vm["main"](input_tensor)[0]
dev.sync()
prediction = output.numpy()

1. tvm.runtime.load_module() 加载 .so，relax.VirtualMachine() 在 musa 设备上创建 VM，并注册编译产物中的 main 函数。

2. tvm.runtime.tensor(input_np, dev) 在 MUSA 上分配张量并复制输入；vm["main"](input_tensor) 调用编译后的前向函数。

3. dev.sync() 等待异步 MUSA 任务完成；output.numpy() 将结果复制回 host，供后续解码和绘制使用。

input_np 已完成 letterbox 预处理，布局为 NCHW、dtype 为 float16。YOLOv8 的原始输出为 (1, 84, 8400)：84 表示 4 个 box 参数加 80 个 COCO 类别，8400 是三个尺度的候选框总数。inference.py 随后完成 box 和类别分数解码、逐类别 NMS、坐标还原和画框。

4. 纯 forward 性能

性能测试只测量 device-resident 前向，不包含图片读取、letterbox、H2D、NMS、D2H 和画框：

for _ in range(10):
    vm["main"](input_tensor)
dev.sync()

evaluator = vm.time_evaluator(
    "main",
    dev,
    number=1,
    repeat=20,
    min_repeat_ms=100,
)
times_ms = np.asarray(evaluator(input_tensor).results) * 1e3

1. 先执行 10 次 vm["main"] 并调用 dev.sync()，完成 runtime 和 kernel 预热。

2. vm.time_evaluator() 以 number=1、repeat=20 采样 20 组；当单组耗时低于 min_repeat_ms=100 时，TVM 会自动增加每组执行次数。

3. evaluator() 返回以秒为单位的计时结果，乘以 1e3 转换为毫秒；latency 取 20 组均值，throughput 按 1000 / latency(ms) 计算。

该口径只反映 TVM-MUSA 模块在 M1000 上的单次前向性能。

Step4 常用参数

构建阶段：

参数	含义	默认值
MODEL	Ultralytics 模型名，例如 yolov8s.pt	yolov8s.pt
--output	输出 .so 路径	artifacts/<model>_musa.so
--imgsz	静态正方形输入尺寸	640

运行时阶段：

参数	含义	默认值
--model	由 build.py 生成的 .so	artifacts/yolov8s_musa.so
--image	实际推理图片	assets/bus.jpg

inference.py 顶部的 WARMUP、REPEAT、CONF_THRESHOLD、IOU_THRESHOLD 和 MAX_DETECTIONS 是固定示例参数，需要调整时直接修改常量。

更换输出路径：

python build.py \
  --output /path/to/yolov8s_musa.so

要更换模型，直接把模型名作为位置参数传入；也可以传本地权重路径：

python build.py yolov8n.pt
python build.py /path/to/custom.pt --output /path/to/custom_musa.so

更换推理图片：

python inference.py \
  --model /path/to/yolov8s_musa.so \
  --image /path/to/demo.jpg

build.py 会在 --output 旁边生成同名 .json；移动 .so 时要一并移动该文件。当前示例针对 YOLOv8 检测头、单输入和静态 shape；换成其他模型时需要相应修改导出方式和后处理。权重和示例图片的版权、许可归各自上游项目所有，对外分发或商用前请确认对应许可条款。