姜汉：使用 TVM 做边缘设备的模型编译和推理实践

关注

原创最新推荐文章于 2026-04-26 14:41:20 发布 1.2k 阅读 AI 写同款· GEO 优化 ›

TVM优化Deep Learning GPU算子 TVM优化Deep Learning GPU算子高效的深度学习算子是深度学习系统的核心。通常，这些算子很难优化，需要HPC专家付出巨大的努力。端到端张量IR / DSL堆栈TVM使这一过程变得更加容易。如何在TVM的帮助下编写高性能GPU运算符内核。本文以深度卷积（即topi.nn.depthwise_conv2d_nchw）为例，并演示如何在tensorflow中改进已经手工优化的CUDA内核。在不同的工作负载下，最终版本比tf-1.2中优化的内核快2到4倍，在启用了算子融合的情况下，最终版本快3到阅读详情

By 超神经

内容一览：本文整理自 OpenBayes贝式计算工程副总裁姜汉在 2023 Meet TVM 上海站的分享，主题为「使用 TVM 做边缘设备的模型编译和推理实践」，HyperAI超神经做了不改变原意的删减整理。

关键词：2023 Meet TVM 机器学习编译

我是来自 OpenBayes贝式计算的姜汉，我的工作是进行 AI 算法开发和做编译器模型推理。今天的分享从使用者的角度，介绍如何用 TVM 做边缘设备的模型编译和推理实践。

OpenBayes贝式计算提供能够在边缘设备上实现图像分类、目标检测、视频追踪的推理模型。这些模型的使用场景会涉及各种不同的硬件，这产生了一个需求，即希望已训练好的模型能在不同设备上跑起来，这符合 TVM 的使用场景。

今天分享的案例，使用的边缘设备是同时搭载 GPU、NPU 的 RK3588 及 RK3399，使用的 TVM 源码是 v0.10 版本。大家可以根据自己的需求选择是否使用交叉编译。

在 Mali (GPU) 上进行模型推理

前情提要

Mali (GPU) 基于 ARM 框架，很多硬件都使用 Mali (GPU)，很多计算库支持 Mali (GPU)。编译过程中软件选择的是 OpenCL 。

用 Mali (GPU) 进行模型推理的编译

首先回顾 TVM 编译流程：

* 导入不同框架的前端模型；

* 生成高层的 RelayIR；

* 生成 TE（张量表示）；

* 使用 Auto-Scheduler 进行自动优化；

* 用训练出的优化方案生成后续的 TE 和 TIR；

* 生成底层中间表示 TIR；

* 编译得到对应目标设备的图标代码，生成可执行文件。

这个过程看起来漫长，写代码时却很容易。直接使用时不会涉及很多优化，可以直接加载模型，确定固定的数据尺寸。

这个案例选择了使用 onnx。在使用 TVM 中，不同的前端框架，onnx 需要确定固定的依赖版本，这样使用时更顺畅，避免使用时因版本不对导致操作不兼容。选择 onnx 还需注意，模型和 TVM 支持的 onnx 算子支持列表对应。

经过编译运行，可以导出直接转化的模型。这个模型没有经过优化，只是翻译了一遍。

使用 Auto Scheduler 优化性能

使用优化涉及到 Auto-Scheduler 的过程：

Auto-Scheduler 作为 TVM 中无模版的优化模块，会在计算中先把计算图拆成子图，再在每一次迭代中进行提取特征、优化、固定、发到指定硬件上评测性能数据指标的操作。Auto-Scheduler 拿到性能指标后，再返回，让系统进行下一步迭代优化，经过一段迭代轮次后，可以达到迭代目标或者完成迭代次数。

Auto-Scheduler 无模版，使用过程中可以不用定义很多参数范围。

如上图，只需定义好网络，把网络放入优化调度中，迭代次数时可设置成 10000 或 20000。模型一般在几个小时内运行完，运行时间会受设备、图片尺寸影响。得到优化模型的可执行文件后，可以加载模型，使用模型处理图像或输入数据，得到更好的性能和识别结果。在 ARM 0 上进行测试，模型效果表现为从普通推理的 200ms 提升到 40ms。

Mali（GPU）是常用设备， TVM 方法有令模型推理少写代码，提升模型性能的效果。

如何使用 NPU 推理模型

案例中实验是在有 NPU 的 RK3588 上进行的，NPU 的支持与 GPU 支持差别巨大，想在 NPU 上推理模型，需要 NPU 提供对应的 API 和算子支持列表，硬件需提供对应的支持和开发。用 NPU 进行推理的实践是在 RK3399 上实现的，RK3588 并不支持用 NPU 进行模型推理。

TVM 需实现对应的后端，即 Codegen。不同的厂商提供了多种硬件，目前 NPU 种类多样。这时使用 TVM 可以分为三种情况：

* TVM 本身支持一种 NPU，这种情况可以直接使用 TVM；

* 厂商提供第三方支持，用自己的库接入 TVM 支持；

* 没有支持，只能选择自己使用开源接口进行后续的开发。

使用 BYOC 方法实践 Codegen，下图为编译流程：

* 生成图的表示，并将其提供给图的引擎，类似 TensorRT（需要第三方支持）；

* 生成汇编，编译器生成可执行文件（需提供动态链接库和 Include 头文件）；

* 如果 NPU 本身没有支持，需按照以下流程进行。

实现 NPU 脚本后，可以将脚本应用在 RK3399 上应用 yolo 进行目标检测实践。这引入了一个问题，如何用 TVM 处理更多的模型呢？这一是需要前端的支持，这样 TVM 模型改动较少，遇到的问题就只有模型本身的动态图问题。

有一种解决方式是直接令动态图尺寸固定，这可能影响模型的性能，但后续不再需要额外操作。如果动态部分只在模型开始或是最后出现，可以直接删掉，工程师在后面接写 CPU 代码。

Relax 社区目前在推进解决动态图问题，但是不提倡一味去等社区的工作成果，可以想办法绕过这类问题。

支持更多模型还涉及到一个问题，即后端的不同支持，这和之前提到的 NPU 提供的支持相似，基于硬件厂商已经开放的设备和接口，做相应的 BYOC 支持。

硬件厂商一般给出非常直接的支持：接口开放、给出算子。工程师只需调用算子即可，并不存在很多可以自己调整的部分。这样，接 BYOC 的优化过程好像只是翻译一遍，优化的效果可能不是非常好，但是解决了把不同框架模型转换到硬件上的问题。

实现应用 TVM 后，工作流程提高了效率，很好地加快了开发速度。

Q&A

Q1：分享一些 TVM 编译步骤在边缘端进行的经验。

A1：主要看硬件怎么搭，如果设备设计的时候就比较完善，像一台小型服务器一样，那么可以直接在上面跑，没必要选择交叉编译。如果系统环境只支持模型驱动，就只能交叉编译。因为盒子已经很小，所以就没必要考虑 Docker 或者其他环境。一般来说，将盒子的环境配成 TVM 编译过程中需要的环境，就和将来执行环境是一套环境，出问题的概率较小。

Q2：能拿到 RK3588 底层的 API 哪些信息？BIOC 这套流程相当于逐段生成对应 API，这样去调 RK3588 硬件，岂不是会在 CPU 和异构设备之间来回倒腾，这种方法的性能相比直接使用 RK3588 提供的 SDK 做 onnx 模型优势在哪？

A2：RK3588 是用不了的，这个问题只能讨论 RK3399。对于 RK3399 有两种办法：

1. 用芯原的 TVM 库。这个库好几年没有用，我将它的代码拿来合并及修 Bug，用它做转载生产代码。

2. 第三方本身支持 TVM ，并需要有一个库。

Q3：将 TVM 部署到具体的 BYOC 硬件遇到什么困难？如果是比较早期的硬件，相应的基础软件不完善的情况下还可以用起来么？使用模拟器？

A3：比较难的部分就是考虑要部署的模型有没有用到它不支持的部分，如果不支持就需要找对应的实现；如果是单纯开发，用模拟器也可以，但我们还是希望能够把模型用起来。因为模拟器是跑在另一个服务器以及非硬件的环境上，只能说在模拟条件能够走通。而且我之前发现能走通模拟器的放到指令硬件上，也有跑不出来的情况。所以我就比较偏向于要么等一等，要么看当下有什么样的条件先跑一下。

Q4：使用 TVM 的考量是什么？

A4：TVM 是一个编译框架。OpenBayes贝式计算本身做算法平台，要提供各种各样的模型，前端现在会有多种框架，后端会有多种设备，硬件也分为 ARM 和 AMD。Openbayes贝式计算平台需要将多样的框架模型转化到多样的硬件的终端下，即有多到多的转化需求，所以选择 TVM。

以上为姜汉在 2023 Meet TVM 上海站的演讲整理内容。接下来，本次活动其他嘉宾分享的详细内容也将陆续在本公众号中发布推出，欢迎持续关注！

获取 PPT：关注微信公众号「HyperAI超神经」，后台回复关键字「TVM 上海」，获取完整 PPT。

更多资料

TVM 中文文档：https://tvm.hyper.ai/

GitHub 地址：https://github.com/apache/tvm

关于 OpenBayes贝式计算

OpenBayes贝式计算是国内领先机器学习及高性能计算服务提供商，通过为新一代异构芯片嫁接经典软件生态及机器学习模型，进而为工业企业及高校科研提供更加快速、易用的数据科学计算产品，其产品已被数十家大型工业场景或头部科研院所所采用。

访问官网：https://openbayes.com/

—— 完 ——