第4章：深度学习框架算子集
各位同学，今天我们来聊聊一个很实际的问题——不同深度学习框架里的算子，到底有什么区别？
说实话，我刚入行那会儿，也被这个问题搞得头大。TensorFlow里一个卷积叫tf.nn.conv2d，PyTorch里叫torch.nn.Conv2d，ONNX里又变成了Conv。名字不一样就算了，参数顺序、默认值、数据排布方式也各有各的脾气。嗯，今天我就把这些坑一个一个给你们讲清楚。
4.1 三大框架的算子生态概览
先看一张总览图，我花了不少时间整理，你们先有个整体印象。
从这张图你能看出来，TensorFlow的算子数量是最多的，PyTorch次之，ONNX最少。为什么？因为ONNX是中间表示格式，它只定义最核心的算子，其他花里胡哨的操作都通过组合来实现。
4.2 算子命名与接口差异
我个人习惯把算子差异分成三类：命名差异、参数差异、行为差异。咱们一个一个说。
4.2.1 命名差异
同一个操作，三个框架叫法完全不同。我整理了一张表，你们直接看：
| 操作类型 | TensorFlow | PyTorch | ONNX |
|---|---|---|---|
| 二维卷积 | tf.nn.conv2d | torch.nn.Conv2d | Conv |
| 最大池化 | tf.nn.max_pool | torch.nn.MaxPool2d | MaxPool |
| 批量归一化 | tf.nn.batch_normalization | torch.nn.BatchNorm2d | BatchNormalization |
| 矩阵乘法 | tf.matmul | torch.mm | MatMul |
| ReLU激活 | tf.nn.relu | torch.nn.ReLU | Relu |
你看，光名字就够你记一阵子的。不过别担心，做NPU算子映射的时候，我们通常以ONNX为基准，因为它是中间表示，最中立。
4.2.2 参数差异——以卷积为例
名字不同还好说，参数差异才是真正的坑。我拿卷积算子举个例子：
// TensorFlow 卷积
// 数据排布: NHWC (默认)
output = tf.nn.conv2d(
input, // [batch, height, width, channels]
filters, // [height, width, in_channels, out_channels]
strides=[1, 2, 2, 1],
padding='SAME',
data_format='NHWC'
)
// PyTorch 卷积
// 数据排布: NCHW (默认)
output = torch.nn.Conv2d(
in_channels=3,
out_channels=64,
kernel_size=3,
stride=2,
padding=1,
bias=True
)
// ONNX 卷积
// 数据排布: NCHW (强制)
output = Conv(
input, // [N, C, H, W]
weight, // [M, C, kH, kW]
bias,
kernel_shape=[3, 3],
strides=[2, 2],
pads=[1, 1, 1, 1],
group=1
)
关键差异点：
- 数据排布：TensorFlow默认NHWC，PyTorch和ONNX默认NCHW。这个差异在NPU上特别要命，因为硬件通常对NCHW更友好。
- 参数传递方式：TensorFlow用位置参数+关键字，PyTorch用类封装，ONNX用属性列表。
- padding策略：TensorFlow的'SAME'和'VALID'与PyTorch的padding数值计算方式不同，映射时要小心。
4.3 算子行为差异——那些年我踩过的坑
我曾经在做一个模型迁移项目时，发现同样的网络结构，TensorFlow和PyTorch推理结果差了0.5%。排查了两天，最后发现是平均池化的边界处理方式不同。
4.3.1 池化算子的边界处理
TensorFlow的tf.nn.avg_pool在计算平均值时，分母是实际参与计算的元素个数，而不是池化窗口大小。而PyTorch的torch.nn.AvgPool2d默认分母是窗口大小，即使有些元素在边界外。
// 举个例子：输入 3x3，池化窗口 2x2，stride=1
// TensorFlow 计算方式：
// 右下角窗口实际只有1个元素在边界内
// avg = 该元素值 / 1
// PyTorch 计算方式：
// 右下角窗口同样只有1个元素
// avg = 该元素值 / 4 (窗口大小)
⚠️ 避坑指南：
我曾经在NPU上实现池化算子时，默认按照PyTorch的行为来写，结果TensorFlow模型跑出来结果不对。后来我加了一个count_include_pad参数来控制这个行为。你们做算子映射时，一定要把这种边界情况考虑进去。
4.3.2 归一化算子的差异
BatchNorm的差异也很大。TensorFlow的tf.nn.batch_normalization需要手动传入均值和方差，而PyTorch的torch.nn.BatchNorm2d内部维护了running_mean和running_var。
ONNX的BatchNormalization算子则把训练和推理模式分开，通过training_mode属性控制。嗯，这里要注意，ONNX的BatchNorm在推理时默认使用训练好的均值和方差，不会重新计算。
4.4 算子映射实战——从PyTorch到ONNX
说完了理论，咱们来点实际的。我给你们演示一下，怎么把一个PyTorch模型里的算子映射到ONNX。
// PyTorch 模型中的算子
class MyModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv = torch.nn.Conv2d(3, 64, 3, padding=1)
self.bn = torch.nn.BatchNorm2d(64)
self.relu = torch.nn.ReLU()
self.pool = torch.nn.MaxPool2d(2)
def forward(self, x):
x = self.conv(x) // → ONNX: Conv
x = self.bn(x) // → ONNX: BatchNormalization
x = self.relu(x) // → ONNX: Relu
x = self.pool(x) // → ONNX: MaxPool
return x
// 映射到ONNX后的算子属性
// Conv:
// kernel_shape: [3, 3]
// strides: [1, 1]
// pads: [1, 1, 1, 1]
// group: 1
// dilations: [1, 1]
// BatchNormalization:
// epsilon: 1e-5
// momentum: 0.9
// training_mode: 0 (推理模式)
// MaxPool:
// kernel_shape: [2, 2]
// strides: [2, 2]
// pads: [0, 0, 0, 0]
💡 我的经验：
做算子映射时，我建议你们先导出ONNX模型，然后用Netron可视化工具看一下。这样能直观地看到每个算子的属性值，比对着文档猜要靠谱得多。
4.5 框架算子差异总结
最后，我把常见的算子差异点总结成一张表，方便你们查阅：
| 算子类别 | 差异点 | 影响范围 | 映射建议 |
|---|---|---|---|
| 卷积 | 数据排布、padding计算、dilation默认值 | 高 | 统一转NCHW，padding按ONNX规范 |
| 池化 | 边界处理、ceil_mode、count_include_pad | 中 | 明确边界策略，测试极端尺寸 |
| 归一化 | 训练/推理模式、epsilon位置、动量计算 | 高 | 推理时固定参数，训练时注意梯度流 |
| 激活函数 | inplace参数、LeakyReLU的alpha默认值 | 低 | 大部分可直接映射，注意参数名 |
| 矩阵运算 | 转置行为、广播规则 | 中 | 明确shape，避免隐式广播 |
| 张量操作 | reshape vs view、permute vs transpose | 中 | 优先用reshape，避免内存不连续 |
说实话，算子映射这件事，没有捷径。你只能一个一个算子去对，去测试。但只要你把常见的那几十个核心算子摸透了，剩下的都是排列组合。
好了，这一章的内容就到这儿。记住，做NPU算子开发，细节决定成败。一个padding算错，整个模型就废了。下节课我们聊聊算子融合，那又是另一个有意思的话题。