TVM Model Loading

基于 Apache TVM 的模型加载与优化示例项目，演示如何使用 TVM Relax 前端导入 ONNX 模型并进行算子融合等优化。

项目结构

.
├── src/
│   ├── main.py              # 入口脚本
│   ├── import_model.py      # 构建自定义 ONNX 模型并导入 TVM Relax
│   └── optimize_model.py    # 模型优化流程（算子计数、标签、融合等）
├── custom_model.onnx        # 生成的示例 ONNX 模型
├── requirements.txt         # Python 依赖
└── README.md

依赖

• Python 3.8+
• Apache TVM
• ONNX
• NumPy

安装依赖：

pip install -r requirements.txt

使用方法

cd src
python main.py

运行后将依次执行：

1. 构建自定义 ONNX 模型 — 创建一个简单的 CNN（Conv → ReLU → MaxPool → Flatten → Gemm）
2. 导入 TVM Relax — 通过 relax.frontend.onnx.from_onnx 将 ONNX 模型转换为 TVM Relax IRModule
3. 模型优化 — 执行自定义 Sequential Pass：
   • CountOpsPass — 统计算子数量
   • TagModelPass — 为模型添加自定义标签
   • LegalizeOps — 将 Relax IR 转换为 TIR
   • AnnotateTIROpPattern — 为 TIR 算子添加融合注解
   • FoldConstant — 常量折叠
   • FuseOps(fuse_opt_level=3) — 算子融合（如 conv+relu+pool → fused_conv_relu_pool）
   • FuseTIR — 合并融合算子的 TIR PrimFunc

自定义模型结构

示例模型为一个简单的图像分类 CNN：

层	输入形状	输出形状
Conv (3×3, 16 filters)	(1, 3, 224, 224)	(1, 16, 224, 224)
ReLU	(1, 16, 224, 224)	(1, 16, 224, 224)
MaxPool (2×2)	(1, 16, 224, 224)	(1, 16, 112, 112)
Flatten	(1, 16, 112, 112)	(1, 200704)
Gemm (→ 10 classes)	(1, 200704)	(1, 10)

许可证

本项目仅供学习和研究使用。