3. 主流AI框架算子集解析：PyTorch、TensorFlow、ONNX概览
做算子映射这些年，我接触最多的就是这三个框架的算子集。说实话，刚开始我也被它们搞得头大——同一个功能，三个框架叫法完全不同，参数顺序也各有脾气。今天咱们就捋一捋，把它们的家底摸清楚。
3.1 PyTorch算子集：动态图的宠儿
PyTorch的算子，我个人觉得是最"人性化"的。为什么？因为它跟Python的思维习惯高度一致。你想想看，一个torch.nn.Conv2d，参数名一看就懂：in_channels、out_channels、kernel_size，基本不用查文档。
我在项目中遇到过最典型的场景：客户给了一个用PyTorch训练的ResNet模型，要求映射到我们的NPU上。PyTorch的算子大概分这么几类：
| 类别 | 典型算子 | 映射难度 |
|---|---|---|
| 张量操作 | torch.view, torch.reshape, torch.permute | 低 |
| 数学运算 | torch.add, torch.mul, torch.matmul | 低 |
| 神经网络层 | torch.nn.Conv2d, torch.nn.Linear | 中 |
| 激活函数 | torch.nn.ReLU, torch.nn.Sigmoid | 低 |
| 归一化 | torch.nn.BatchNorm2d, torch.nn.LayerNorm | 中 |
| 池化 | torch.nn.MaxPool2d, torch.nn.AvgPool2d | 低 |
嗯，这里要注意：PyTorch的torch.view和torch.reshape虽然功能相似，但view要求张量在内存中是连续的。我踩过这个坑——模型在GPU上跑得好好的，一上NPU就报错，查了半天发现是view之前有个permute操作，内存不连续了。
核心要点：PyTorch算子映射时，重点关注torch.nn.functional下的函数式接口，它们往往比torch.nn.Module更底层，映射起来更灵活。
3.2 TensorFlow算子集：静态图的坚守者
TensorFlow的算子，说白了就是"规矩多"。它分两套API：tf.*和tf.keras.*。我个人习惯用tf.keras，因为跟PyTorch的nn.Module风格更接近。
TensorFlow算子有个特点：很多操作都带name参数。这其实是个好习惯——调试的时候，你能从错误信息里直接定位到是哪个节点出了问题。我曾经帮一个团队排查NPU推理错误，就是靠算子名称反推回原始模型结构，半小时就定位到了问题。
常见的TensorFlow算子分类：
- 张量操作：tf.reshape, tf.transpose, tf.concat
- 数学运算：tf.add, tf.multiply, tf.matmul
- 卷积相关：tf.nn.conv2d, tf.nn.depthwise_conv2d
- 池化：tf.nn.max_pool2d, tf.nn.avg_pool2d
- 激活：tf.nn.relu, tf.nn.sigmoid, tf.nn.softmax
避坑指南：我曾经把TensorFlow的tf.nn.conv2d直接按PyTorch的Conv2d逻辑映射，结果精度对不上。后来发现，TensorFlow默认的data_format是"NHWC"，而PyTorch是"NCHW"。这个差异，你映射的时候一定要检查。
3.3 ONNX算子集：中间人的智慧
ONNX，说白了就是AI框架界的"通用语言"。它不偏向任何框架，只定义了一套标准算子。我特别喜欢ONNX的一点：它的算子定义非常清晰，每个算子的输入输出、属性、形状推导规则都写得明明白白。
ONNX算子集目前有150多个算子（opsets版本不同数量有差异），但常用的也就30-40个。核心算子包括：
| 算子名 | 功能 | 对应PyTorch | 对应TensorFlow |
|---|---|---|---|
| Conv | 卷积 | torch.nn.Conv2d | tf.nn.conv2d |
| Gemm | 全连接 | torch.nn.Linear | tf.keras.layers.Dense |
| Relu | ReLU激活 | torch.nn.ReLU | tf.nn.relu |
| MaxPool | 最大池化 | torch.nn.MaxPool2d | tf.nn.max_pool2d |
| BatchNormalization | 批归一化 | torch.nn.BatchNorm2d | tf.keras.layers.BatchNormalization |
| Reshape | 形状变换 | torch.reshape | tf.reshape |
你想想看，ONNX最大的价值是什么？是标准化。不管你的模型是用PyTorch训练的，还是TensorFlow训练的，只要导出成ONNX格式，NPU的算子映射就只需要针对ONNX做一套。这大大降低了开发工作量。
个人经验：我建议你在做算子映射时，优先考虑ONNX作为中间表示。如果客户给的是PyTorch模型，先转ONNX再映射；如果是TensorFlow模型，同样先转ONNX。这样你只需要维护一套ONNX到NPU的映射规则，省心很多。
3.4 三个框架的算子映射对比
我整理了一张对比表，方便你快速查阅：
| 特性 | PyTorch | TensorFlow | ONNX |
|---|---|---|---|
| 动态/静态图 | 动态图（Eager） | 静态图（Graph） | 静态图 |
| 算子数量 | 约200+ | 约1000+（含tf.keras） | 约150+ |
| 数据格式 | NCHW（默认） | NHWC（默认） | NCHW（推荐） |
| 自定义算子 | 容易（继承nn.Module） | 较复杂（需注册） | 需扩展opsets |
| 调试难度 | 低（可逐行调试） | 中（需构建图） | 中（需可视化工具） |
| 映射复杂度 | 中 | 高 | 低 |
为什么会这样？说白了，PyTorch的动态图让你能随时打印张量值，调试起来特别方便。TensorFlow的静态图虽然优化空间大，但出了问题排查起来费劲。ONNX作为中间格式，算子定义最精简，映射起来自然最轻松。
3.5 知识体系总览
下面这张图，是我根据多年经验总结的算子映射知识体系。你可以把它当作一个"导航图"：
这张图展示了算子映射的完整链路：从三个主流框架出发，经过中间表示层，最终落到NPU的算子映射层。我建议你按照这个路径来设计你的映射方案，会少走很多弯路。
总结一下：PyTorch算子最直观，TensorFlow算子最全面，ONNX算子最标准。做NPU算子映射时，我的建议是——以ONNX为中间桥梁，以PyTorch为调试参考，以TensorFlow为兼容性补充。这样三管齐下，基本能覆盖90%以上的实际场景。
想要更多资料，联系可可