Dynamic Model in TVM

• 语言: zh
• 格式: .pdf
• 来源: github.com/tvmai

摘要

文档主要探讨了在TVM框架下支持动态模型的技术挑战与解决方案。TVM作为高性能的深度学习推理框架，其当前的图运行时无法直接编译和运行包含动态控制流、动态输入形状和依赖数据的输出形状等特性的动态模型。为了解决这一问题，文档提出了通过支持任意维度的类型推断、使用形状函数在运行时计算类型以及引入虚拟机作为新的运行时环境来支持动态模型。此外，文档还讨论了动态代码生成的目标、挑战以及实现方案，包括内核调度和子图调度的策略。最后，文档通过API示例展示了如何在ResNet-50等模型中应用这些技术。

AI总结

# 《Dynamic Model in TVM》 总结

## 动态模型的特点
动态模型具有以下特点：
1. **控制流**：支持条件语句（if、else）和循环（while、for）等动态执行路径。
2. **动态形状**：
- 输入形状动态：如批次大小、图像大小、序列长度等。
- 输出形状依赖于数据：如`arange`、`nms`等操作。
- 控制流中的动态拼接：如在`while`循环中拼接张量。

## TVM的局限性
TVM的原始图运行时无法直接支持动态模型：
1. 无法编译和运行包含控制流或动态形状的模型。
2. 对于动态模型的性能较差：
- 单个内核在不同形状上的表现不一致。
- 需要为同一操作生成不同模板。
- TVM的计算与调度耦合，限制了灵活性。

## 支持动态模型的方案
为了支持动态模型，TVM提出了以下改进方向：
1. **支持任意维度的类型**：通过类型推断和运行时形状计算，增强类型系统的灵活性。
2. **虚拟机作为新运行时**：通过虚拟机（VM）实现动态模型的执行，支持动态代码生成。
3. **动态代码生成**：
- **操作分发**：根据操作的条件和形状，动态选择合适的实现。
- **内核分发**：为不同形状生成不同的内核，优化性能。
- **图分发**：在子图级别进行动态分发。

## 动态代码生成的提案
1. **目标**：支持动态形状的代码生成。
2. **挑战**：
- 单个内核在不同形状上的性能差异较大。
- 需要为同一操作生成多个模板。
- TVM的计算与调度耦合，影响灵活性。

## 动态代码生成的实现思路
1. **操作分发**：
- 通过`OpStrategy`注册默认实现和专用实现。
- 根据操作的条件和形状动态选择实现。
2. **内核分发**：
- 为不同形状生成不同的内核。
- 使用`SpecializedConditionNode`进行条件判断。
3. **图分发**：
- 在子图级别进行动态分发，优化整体性能。

## 示例代码
以下是一个动态模型的编译和运行示例：

# 定义输入形状
input_name = "data"
input_shape = [tvm.relay.Any(), 3, 224, 224]
dtype = "float32"

# 加载模型
block = get_model('resnet50_v1', pretrained=True)
mod, params = relay.frontend.from_mxnet(block, shape={input_name: input_shape}, dtype=dtype)

# 注册动态分发策略
tvm.relay.transform.dispatch_global_func(mod, "main", {input_name: input_shape}, tvm.relay.vm.exp_dispatcher)

# 编译模型
vmc = relay.backend.vm.VMCompiler()
with tvm.autotvm.apply_graph_best("resnet50_v1_graph_opt.log"):
    vm = vmc.compile(mod, "llvm")

# 初始化和加载参数
vm.init(ctx)
vm.load_params(params)

# 运行测试
data = np.random.uniform(size=(1, 3, 224, 224)).astype("float32")
out = vm.run(data)

data = np.random.uniform(size=(4, 3, 224, 224)).astype("float32")
out = vm.run(data)

## 总结
1. 动态模型的特点：控制流和动态形状使得模型更加灵活，但也带来了编译和运行时的挑战。
2. TVM的局限性：无法直接支持动态模型，性能和灵活性不足。
3. 支持动态模型的方案：通过虚拟机、动态代码生成和操作分发，增强TVM对动态模型的支持。
4. 动态代码生成的实现思路：基于操作和内核的分发策略，优化性能和灵活性。

未来的工作方向包括优化动态代码生成的效率，进一步完善虚拟机的运行时支持，以及扩展对更多动态模型的兼容性。