2. 计算图与算子表示：DAG基础、框架中的算子、ONNX中间表示

好，咱们正式开始第二讲。这一章我打算聊聊计算图——说白了，就是AI模型在底层到底长什么样。你写几行PyTorch代码，框架背后帮你搭了一张图，然后编译器再对着这张图做优化。不理解这张图，后面谈算子映射和融合，基本就是空中楼阁。

我个人习惯把计算图比作一张「菜谱」。算子就是一个个烹饪步骤，数据就是食材，依赖关系就是「先切菜再炒菜」。你想想看，编译器拿到这张菜谱，才能决定哪些步骤可以合并，哪些步骤可以并行。

2.1 计算图（DAG）基础

计算图本质上是一个有向无环图（DAG）。节点是算子（比如卷积、ReLU、矩阵乘），边是数据流（张量）。为什么必须是无环？因为推理过程是前向传播，数据不能走回头路。如果有环，那就变成循环神经网络那种带状态的图了，处理方式完全不同。

我在项目中遇到过一个问题：有同事把模型导出时，发现图里莫名其妙多了一个环。查了半天，原来是某个自定义算子内部引用了自身的输出。嗯，这里要注意——编译器遇到环会直接报错，因为没法做静态调度。

核心要点：

  • 节点：算子（Conv2D、MatMul、Add、ReLU...）
  • 边：张量（Tensor），携带形状、数据类型、内存地址
  • 拓扑序：编译器按依赖关系排序，决定执行顺序

下面这张图，是我手绘的一个简单计算图结构。你可以看到数据从输入开始，经过两个分支，最后汇合输出。

Input (Tensor) Conv2D ReLU Add Output 分支1 分支2

这张图里，Conv2D和ReLU是并行执行的——它们都依赖同一个输入。编译器看到这种结构，就会考虑把两个算子调度到不同的计算单元上。这就是后面要讲的「并行优化」的基础。

我的经验：实际模型的计算图往往有几百甚至上千个节点。调试时别直接看全图，先看子图。我习惯用Netron工具可视化，一眼就能看出哪些分支是瓶颈。

2.2 TensorFlow/PyTorch中的算子表示

不同框架对算子的表示方式不一样。但底层逻辑是相通的——都是把算子抽象成一个对象，包含输入、输出、属性和类型。

2.2.1 TensorFlow的算子表示

TF早期用的是静态图（Graph模式），后来加了Eager模式。但底层表示一直是tf.Operation和tf.Tensor。每个算子就是一个Operation，输入输出都是Tensor。

# TensorFlow 静态图示例
import tensorflow as tf

# 定义占位符
a = tf.placeholder(tf.float32, shape=[None, 32], name='input_a')
b = tf.placeholder(tf.float32, shape=[None, 32], name='input_b')

# 算子：MatMul + Add
c = tf.matmul(a, b, name='matmul_op')
d = tf.add(c, 1.0, name='add_op')

# 查看图结构
print(d.op.name)        # add_op
print(d.op.inputs)      # [<tf.Tensor 'matmul_op:0' shape=(?, 32) dtype=float32>, <tf.Tensor 'Const:0' shape=() dtype=float32>]
print(d.op.node_def)    # 序列化的NodeDef协议

你看，每个算子都有一个node_def，里面记录了算子类型、属性、输入输出名称。编译器就是靠这个来解析图的。

我曾经踩过一个坑：TF的静态图在导出时，如果某个算子的输入张量形状是动态的（比如None），某些后端编译器会直接拒绝编译。所以导出前最好用tf.TensorShape显式指定形状。

2.2.2 PyTorch的算子表示

PyTorch用的是动态图（Define-by-Run）。每次前向传播都会重新构建图。但别被「动态」两个字骗了——底层依然有图结构，只是每次执行时动态生成。

# PyTorch 动态图示例
import torch

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = torch.nn.Conv2d(3, 64, 3)
        self.relu = torch.nn.ReLU()
    
    def forward(self, x):
        x = self.conv(x)
        x = self.relu(x)
        return x

model = MyModel()
dummy_input = torch.randn(1, 3, 224, 224)

# 使用 torch.jit.trace 导出为静态图
traced_model = torch.jit.trace(model, dummy_input)
print(traced_model.graph)
# graph(%self : __torch__.MyModel,
#       %input : Float(1:150528, 3:50176, 224:224, 224:1)):
#   %conv = prim::GetAttr[name="conv"](%self)
#   %relu = prim::GetAttr[name="relu"](%self)
#   %1 = prim::CallMethod[name="forward"](%conv, %input)
#   %2 = prim::CallMethod[name="forward"](%relu, %1)
#   return (%2)

PyTorch的图表示是torch::jit::Graph，节点是Node，值是Value。每个Node包含算子类型（比如aten::conv2d）、输入输出Value列表。

注意：PyTorch的JIT trace有个坑——它只能trace不包含控制流的路径。如果你的模型里有if或for循环，trace出来的图可能不完整。这时候得用torch.jit.script，但script对Python语法的支持有限。

2.3 ONNX中间表示

ONNX（Open Neural Network Exchange）是微软和Facebook搞的「通用语言」。说白了，就是让不同框架的模型能互相转换。你想想看，如果每个芯片厂商都要适配TF、PyTorch、PaddlePaddle……那不得累死？有了ONNX，大家只要对接ONNX就行了。

2.3.1 ONNX的算子集

ONNX定义了一套标准算子，比如Conv、Relu、Gemm（通用矩阵乘）。每个算子都有明确的输入输出规范和属性定义。

算子名称 输入 输出 关键属性
Conv X, W, B (可选) Y kernel_shape, strides, pads, dilations
Relu X Y 无
Gemm A, B, C Y transA, transB, alpha, beta
BatchNormalization X, scale, B, mean, var Y, running_mean, running_var epsilon, momentum

每个算子还有一个版本号。比如Conv在opset 1和opset 11里，属性定义可能不一样。我建议导出ONNX时指定一个稳定的opset版本（比如11或13），别用最新的——新版本可能有些芯片编译器还没支持。

2.3.2 ONNX的图结构

ONNX的图结构用Protobuf序列化。核心数据结构是ModelProto -> GraphProto -> NodeProto / TensorProto / ValueInfoProto。

// ONNX 图结构的简化表示（伪代码）
message GraphProto {
    repeated NodeProto node;          // 所有算子节点
    repeated TensorProto initializer; // 常量（权重、偏置）
    repeated ValueInfoProto input;    // 图输入
    repeated ValueInfoProto output;   // 图输出
    string name;
}

message NodeProto {
    string op_type;      // 算子类型，如 "Conv"
    repeated string input;   // 输入张量名称列表
    repeated string output;  // 输出张量名称列表
    repeated AttributeProto attribute; // 算子属性
    string name;         // 节点名称（可选）
}

你看，每个节点只记录输入输出的名称字符串，不直接存张量数据。张量数据放在initializer里。这种设计让图结构非常轻量，方便传输和解析。

关键点：ONNX的图是静态的、显式的。所有控制流（比如if、loop）都用专门的算子（If、Loop）表示，而不是Python的if语句。这意味着你导出的ONNX模型，没有隐藏的分支，编译器可以放心优化。

2.3.3 从PyTorch导出ONNX的实战

我直接给你看一段我常用的导出代码。这里面有几个坑，我一个个说。

import torch
import torch.onnx

# 假设你有一个训练好的模型
model = MyModel()
model.eval()  # 一定要切换到eval模式！否则BN和Dropout的行为会不同

dummy_input = torch.randn(1, 3, 224, 224)

torch.onnx.export(
    model,                     # 模型
    dummy_input,               # 示例输入（用于trace）
    "model.onnx",              # 输出路径
    export_params=True,        # 导出权重参数
    opset_version=13,          # 我习惯用13，兼容性好
    do_constant_folding=True,  # 常量折叠优化
    input_names=['input'],     # 输入名称
    output_names=['output'],   # 输出名称
    dynamic_axes={
        'input': {0: 'batch_size'},   # 动态batch
        'output': {0: 'batch_size'}
    }
)

# 验证导出的ONNX模型
import onnx
onnx_model = onnx.load("model.onnx")
onnx.checker.check_model(onnx_model)
print("ONNX模型验证通过！")

这里有几个我踩过的坑：

  • eval模式：忘记切eval模式，导出的BN层会带着训练时的统计量，推理结果直接崩。
  • dynamic_axes：如果你的模型需要支持不同batch size，一定要设置这个。否则导出的模型固定batch=1。
  • do_constant_folding：建议开启。它会把一些常量计算（比如1+2）在导出时就算好，减少运行时开销。

我的习惯：导出ONNX后，我会用onnxruntime跑一遍推理，跟PyTorch的结果对比。误差超过1e-5就要查原因。很多时候是算子精度问题，比如某些芯片的LayerNorm实现跟PyTorch不完全一致。

2.4 小结

这一章我们聊了计算图的基础——DAG结构、TF和PyTorch的算子表示、ONNX中间格式。说白了，计算图就是AI模型的「骨架」。你理解了骨架，后面做算子映射和融合，才知道哪些地方能下手优化。

我个人觉得，ONNX是目前最实用的中间表示。它把不同框架的模型统一成一个标准格式，让芯片厂商只需要对接ONNX就行。但要注意，ONNX也不是万能的——有些框架特有的算子（比如TF的某些高级API）可能找不到对应的ONNX算子，这时候就得自己写自定义算子。

好，这一章就到这里。下一章我们聊聊算子映射——怎么把ONNX里的算子，映射到芯片的指令集上。那才是真正考验编译器功底的地方。

避坑指南：我曾经在导出ONNX时，发现模型里有个torch.nn.Upsample，ONNX里对应的Resize算子在不同opset版本下行为不一样。opset 10用nearest模式，opset 11改成了coordinate_transformation_mode。如果你不指定，默认行为可能跟预期不同。所以导出前一定要查一下目标opset的算子规范。

想要更多资料，联系可可

← 上一章 📖 返回目录 下一章 →