面向机器学习模型的图抽象（Graph Abstraction）

原文：Graph Abstraction for ML Models — Apache TVM

目录

• 什么是图抽象？
• Relax 的关键特性

面向机器学习模型的图抽象

图抽象（graph abstraction）是机器学习（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 所采用的图表示，通过若干关键特性支撑 ML 模型的端到端优化：

• 一等公民式符号形状（First-class symbolic shape）：Relax 使用符号形状表示张量维度，从而能够在张量算子与函数调用之间全局跟踪动态形状关系。

• 多层抽象（Multi-level abstractions）：Relax 支持跨层抽象，从高层神经网络层到低层张量运算均可表达，从而支持跨越模型不同层次的优化。

• 可组合变换（Composable transformations）：Relax 提供可组合变换框架，可选择性地作用于模型的不同组成部分。这包括部分降级（partial lowering）与部分特化（partial specialization）等能力，为定制与优化提供灵活选项。

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