ML 模型的图抽象

图抽象是机器学习 (ML) 编译器中使用的关键技术，用于表示和推理 ML 模型的结构和数据流。通过将模型抽象为图表示，编译器可以执行各种优化以提高性能和效率。本教程将介绍图抽象的基础知识、Relax IR 的关键要素，以及它如何在 ML 编译器中实现优化。

什么是图抽象？

图抽象是将 ML 模型表示为有向图的过程，其中节点表示计算操作（例如，矩阵乘法、卷积），边表示这些操作之间的数据流。这种抽象允许编译器分析模型不同部分之间的依赖关系和关联。

from tvm.script import relax as R

@R.function
def main(
    x: R.Tensor((1, 784), dtype="float32"),
    weight: R.Tensor((784, 256), dtype="float32"),
    bias: R.Tensor((256,), dtype="float32"),
) -> R.Tensor((1, 256), dtype="float32"):
    with R.dataflow():
        lv0 = R.matmul(x, weight)
        lv1 = R.add(lv0, bias)
        gv = R.nn.relu(lv1)
        R.output(gv)
    return gv

Relax 的主要特性

Relax，Apache TVM Unity 策略中使用的图表示，通过几个关键特性促进 ML 模型的端到端优化

• 一流的符号形状：Relax 采用符号形状来表示张量维度，从而能够全局跟踪张量运算符和函数调用之间的动态形状关系。

• 多级抽象：Relax 支持跨级抽象，从高层神经网络层到低层张量操作，从而实现跨模型内不同层次结构的优化。

• 可组合的转换：Relax 提供了一个用于可组合转换的框架，可以有选择地应用于不同的模型组件。这包括诸如部分降级和部分特化等功能，提供灵活的自定义和优化选项。

这些特性共同使 Relax 能够在 Apache TVM 生态系统中为 ML 模型优化提供强大且适应性强的方法。