Apache TVM 快速入门与集成示例

Apache TVM 快速入门实操教程：从 import 验证到 TVMScript 构图、llvm 编译、VM 执行，再到 ONNX 模型端到端示例，全部代码附实测运行回显。

Apache TVMApache TVM2026/10/12

Apache TVM 快速入门与集成示例

本文基于 Apache TVM 0.27.0.post1（Ubuntu 22.04 / Python 3.10 与 CentOS 8 / Python 3.11 双环境实测），提供两条可直接复制的上手路径：不依赖任何外部模型的 TVMScript 最小示例，以及 ONNX 模型端到端示例。安装步骤见《Ubuntu 22.04 安装 Apache TVM 指南》或《CentOS 8 安装 Apache TVM 指南》。

一、核心概念速览

  • IRModule：TVM 中间表示的统一载体，包含一个或多个函数。
  • Relax：0.20 之后主推的图级 IR（Unity 路线），替代早期 Relay 的定位。
  • relax.build：把 IRModule 编译为可执行产物（VMExecutable）。
  • VirtualMachine：执行编译产物的运行时。
  • target：编译目标，CPU 场景用 "llvm"（或 "c"）。

一个值得注意的实测事实：0.27 的运行时张量入口是 tvm.runtime.tensor，旧教程里的 tvm.nd.array 已移除。

二、最小端到端示例（TVMScript 构图）

2.1 完整代码

保存为 tvm_min_demo.py：

"""TVM 最小端到端示例：构造 Relax IRModule -> 编译到 llvm -> VM 执行 -> 与 NumPy 对比"""
import numpy as np
import tvm
from tvm import relax
from tvm.script import relax as R

print("TVM version:", tvm.__version__)

@R.function
def main(x: R.Tensor((2, 3), "float32"), w: R.Tensor((3, 4), "float32")) -> R.Tensor((2, 4), "float32"):
    with R.dataflow():
        gv: R.Tensor((2, 4), "float32") = R.matmul(x, w)
        R.output(gv)
    return gv

mod = tvm.ir.IRModule({"main": main})
target = tvm.target.Target("llvm")
ex = relax.build(mod, target)
dev = tvm.runtime.device("cpu")
vm = relax.VirtualMachine(ex, dev)

x_np = np.random.rand(2, 3).astype("float32")
w_np = np.random.rand(3, 4).astype("float32")
out = vm["main"](tvm.runtime.tensor(x_np, dev), tvm.runtime.tensor(w_np, dev))

print("输出形状:", out.shape, "dtype:", out.dtype)
print("输出前 2 个元素:", out.numpy()[0][:2])
np.testing.assert_allclose(out.numpy(), x_np @ w_np, rtol=1e-4, atol=1e-5)
print("与 NumPy matmul 结果一致: PASS")

2.2 实测运行回显

/opt/tvm-venv/bin/python tvm_min_demo.py

TVM version: 0.27.0.post1
输出形状: (2, 4) dtype: float32
输出前 2 个元素: [0.4101006  0.27392718]
与 NumPy matmul 结果一致: PASS

实测整套流程含 LLVM 编译约 2.5 秒（12 核 x86_64）。代码路径与数据流：@R.function 声明图函数 → relax.build 生成 VMExecutable → VirtualMachine 加载 → 传入 tvm.runtime.tensor 执行 → out.numpy() 取回结果。

2.3 集成要点

  • 输入张量必须与声明 shape/dtype 完全一致，否则在类型检查阶段报错。
  • dev = tvm.runtime.device("cpu")；GPU 场景替换为 tvm.runtime.device("cuda", 0) 并相应调整 target。
  • 产物 ex 可通过 ex.save_lib / export_library 落盘，在部署机直接 tvm.runtime.load_module 加载，实现"编译机与推理机分离"。

三、ONNX 模型端到端示例

3.1 准备依赖

/opt/tvm-venv/bin/pip install onnx

实测安装 onnx 1.23.2。

3.2 完整代码

保存为 tvm_onnx_demo.py。示例用 onnx.helper 构造一个极简 Gemm 模型（避免下载外部权重，便于离线复现）：

"""ONNX 模型端到端示例：onnx.helper 造模型 -> relax 前端导入 -> llvm 编译 -> VM 执行 -> NumPy 对比"""
import numpy as np
import onnx
from onnx import helper, TensorProto
import tvm
from tvm import relax
from tvm.relax.frontend import onnx as onnx_frontend

rng = np.random.default_rng(42)
W_np = rng.standard_normal((8, 16)).astype("float32")
B_np = rng.standard_normal((16,)).astype("float32")

# 1. 构造一个极简 ONNX 模型: y = x @ W + B (Gemm, alpha=beta=1)
X = helper.make_tensor_value_info("x", TensorProto.FLOAT, [4, 8])
Y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [4, 16])
W_init = helper.make_tensor("W", TensorProto.FLOAT, [8, 16], W_np.flatten().tolist())
B_init = helper.make_tensor("B", TensorProto.FLOAT, [16], B_np.flatten().tolist())
node = helper.make_node("Gemm", ["x", "W", "B"], ["y"])
graph = helper.make_graph([node], "tiny_gemm", [X], [Y], initializer=[W_init, B_init])
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
onnx.checker.check_model(model)
onnx.save(model, "tiny_gemm.onnx")
print("已生成 tiny_gemm.onnx")

# 2. 通过 relax ONNX 前端导入
mod = onnx_frontend.from_onnx(onnx.load("tiny_gemm.onnx"), shape_dict={"x": (4, 8)})
print("Relax IRModule 导入成功, 包含函数:", [k for k, _ in mod.functions.items()])

# 3. 编译到 llvm (CPU)
ex = relax.build(mod, target="llvm")
print("编译完成:", type(ex).__name__)

# 4. 在 CPU 上运行
dev = tvm.runtime.device("cpu")
vm = relax.VirtualMachine(ex, dev)
x_np = rng.standard_normal((4, 8)).astype("float32")
y_tvm = vm["main"](tvm.runtime.tensor(x_np, dev))
print("TVM 输出形状:", y_tvm.shape)

# 5. 与 NumPy 参考实现对比
y_ref = x_np @ W_np + B_np
np.testing.assert_allclose(y_tvm.numpy(), y_ref, rtol=1e-4, atol=1e-5)
print("TVM 与 NumPy 参考结果一致: PASS")

3.3 实测运行回显

/opt/tvm-venv/bin/python tvm_onnx_demo.py

已生成 tiny_gemm.onnx
Relax IRModule 导入成功, 包含函数: [I.GlobalVar("main")]
编译完成: VMExecutable
TVM 输出形状: (4, 16)
TVM 与 NumPy 参考结果一致: PASS

Ubuntu 22.04（Python 3.10.12 + onnx 1.23.2）与 CentOS 8（Python 3.11.9 + onnx 1.23.2）双环境实测均 PASS。真实业务模型只需把第 2 步的 onnx.load("tiny_gemm.onnx") 换成自己的模型文件，并在 shape_dict 中给出输入尺寸。

3.4 集成到业务系统的典型形态

# 编译一次
ex = relax.build(mod, target="llvm")
# 生产环境常驻：VirtualMachine 随服务进程加载
vm = relax.VirtualMachine(ex, tvm.runtime.device("cpu"))
# 每次请求
y = vm["main"](tvm.runtime.tensor(x_batch, tvm.runtime.device("cpu")))

四、常见问题

**Q1：**报 AttributeError: module 'tvm' has no attribute 'nd' 或 tvm.nd.array 不存在？ 0.27 已移除 tvm.nd，改用 tvm.runtime.tensor(numpy_array, dev)。

**Q2：**报 from_onnx 接受 str 却出现 'str' object has no attribute 'ir_version'？ 0.27 的 tvm.relax.frontend.onnx.from_onnx() 第一参数要求 onnx.ModelProto 对象，不再接受文件路径字符串。先 onnx.load("model.onnx") 再传入（实测 0.27 行为）。

**Q3：**报 te.create_schedule 不存在？ 经典 TE 调度接口在 0.27 已删除，图级编程请用 Relax；确需 TE 时使用 te.create_prim_func 配合新构建流程。

**Q4：**target 写 "llvm -mcpu=core-avx2" 报错？ 0.27 移除了 CLI 字符串形式，改用 tvm.target.Target({"kind": "llvm", "mcpu": "core-avx2"})。

Q5：import tvm 报 No module named 'tornado'？ 只有 RPC tracker 等可选功能需要 tornado，安装 pip install "apache-tvm[rpc]"；本篇示例不依赖它。

**Q6：**想要命令行工具 tvmc？ 0.27 的 PyPI 包不含 tvmc（0.20.0 起移除）。用法与替代方案见《Apache TVM 常用 CLI 命令与工具详解》。

五、参考链接