3. 主流AI框架介绍：TensorFlow、PyTorch、MindSpore的模型表示与中间表示（IR）

做NPU算子开发这些年，我接触过不少AI框架。说实话，每个框架都有自己的脾气。但不管它们怎么变，核心就两件事：怎么描述模型，以及怎么把模型转成中间表示（IR）。今天我就结合自己的踩坑经历，聊聊这三个主流框架。

3.1 TensorFlow：静态图的坚守者

TensorFlow的模型表示，说白了就是一张静态计算图。你先把所有操作搭好，然后塞数据进去跑。我刚开始接触时特别不习惯——明明写个for循环就能搞定的事，在TF里得用tf.while_loop。

核心概念：Graph + Session。图是静态的，Session负责执行。

举个例子，一个简单的线性回归：

import tensorflow as tf

# 构建图
x = tf.placeholder(tf.float32, shape=[None, 1])
w = tf.Variable(tf.random_normal([1, 1]))
b = tf.Variable(tf.zeros([1]))
y = tf.matmul(x, w) + b

# 损失函数
loss = tf.reduce_mean(tf.square(y - tf.placeholder(tf.float32)))

# 训练
train_op = tf.train.GradientDescentOptimizer(0.01).minimize(loss)

# 执行
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    for step in range(100):
        sess.run(train_op, feed_dict={x: data, y: labels})

看到没？图是提前定义好的，数据通过feed_dict往里灌。这种模式的好处是——优化空间大。NPU在编译时就能把整个图看透，做算子融合、内存复用。我在做某款芯片的TF算子适配时，就靠这个特性把推理性能提了30%。

TF的中间表示是GraphDef（protobuf格式），后来进化到MLIR。GraphDef说白了就是个序列化的图结构，节点是Op，边是Tensor。你想想看，NPU编译器拿到这个，直接解析就能知道整个计算流程。

我的经验：做TF算子映射时，重点关注tf.raw_ops里的底层算子。高层API（如Keras）最终也会拆成这些原子操作。

3.2 PyTorch：动态图的灵活派

PyTorch走的是另一条路——动态图（Define-by-Run）。你写代码的同时，图就建好了。我个人特别喜欢这种风格，调试起来太方便了。

同样的线性回归，PyTorch写法：

import torch
import torch.nn as nn

class LinearRegression(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(1, 1)
    
    def forward(self, x):
        return self.linear(x)

model = LinearRegression()
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

for epoch in range(100):
    optimizer.zero_grad()
    outputs = model(data)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()

动态图的好处是灵活。我在做某个NLP模型的NPU适配时，模型里有条件分支和循环，PyTorch直接写Python逻辑就行，TF就得折腾半天。但代价是什么？性能优化难。因为图是动态生成的，编译器没法提前做全局优化。

PyTorch的中间表示是TorchScript（通过torch.jit.trace或torch.jit.script得到）。它会把你写的Python代码转成静态图，序列化成.pt文件。嗯，这里要注意：trace只能记录实际执行路径，遇到控制流会出问题；script能处理控制流，但有些Python特性不支持。

我曾经踩过的坑：用torch.jit.trace导出一个带if-else的模型，结果推理时永远走同一个分支。后来改用script才解决。所以，如果你的模型有控制流，老老实实用script。

3.3 MindSpore：全场景的后来者

MindSpore是华为推出的框架，主打全场景（端、边、云）。它的模型表示融合了静态图和动态图的优点——动静统一。你写代码时像PyTorch一样动态，但框架会自动转成静态图执行。

看个例子：

import mindspore as ms
from mindspore import nn, Tensor

class LinearRegression(nn.Cell):
    def __init__(self):
        super().__init__()
        self.linear = nn.Dense(1, 1)
    
    def construct(self, x):
        return self.linear(x)

model = LinearRegression()
loss_fn = nn.MSELoss()
optimizer = nn.SGD(model.trainable_params(), learning_rate=0.01)

def forward_fn(x, y):
    pred = model(x)
    return loss_fn(pred, y)

grad_fn = ms.value_and_grad(forward_fn, None, optimizer.parameters)

for epoch in range(100):
    loss, grads = grad_fn(data, labels)
    optimizer(grads)

注意construct方法——它定义了计算逻辑。MindSpore会分析这个函数，自动生成静态计算图。我在做算子开发时，发现MindSpore的MindIR（中间表示）设计得很清晰。它基于图表示，支持函数式语义，而且能保留源码级别的调试信息。

MindIR的特点：

基于ANF（A-Normal Form）中间表示，每个表达式都绑定到变量
支持子图嵌套，方便做算子融合
保留类型信息，NPU编译器可以直接利用

3.4 三者的IR对比

我整理了一张对比表，方便你快速了解：

特性	TensorFlow	PyTorch	MindSpore
图类型	静态图	动态图	动静统一
IR格式	GraphDef / MLIR	TorchScript	MindIR
序列化	.pb / .mlir	.pt	.mindir
控制流支持	通过tf.cond/tf.while	原生Python控制流	自动转成子图
NPU友好度	高（静态图优化空间大）	中（需先转TorchScript）	高（原生支持昇腾）

你想想看，做NPU算子映射时，最头疼的是什么？是IR解析。每个框架的IR格式不同，你得写不同的解析器。我建议团队统一用ONNX作为中间桥梁——大部分框架都支持导出ONNX，然后NPU编译器只对接ONNX就行。

3.5 从IR到NPU算子的映射流程

不管哪个框架，最终都要走到这一步：IR → 算子识别 → 硬件映射。我画了个流程图，帮你理清思路：

这个流程我走了无数遍。每次拿到一个新框架的模型，第一步就是看它的IR结构。比如TF的GraphDef，节点里有个op字段标明算子类型，attr字段存属性。PyTorch的TorchScript则是扁平的指令列表，每条指令有kind和inputs。MindSpore的MindIR更接近函数式，每个节点都是纯函数调用。

我的建议：刚开始做算子映射时，先挑一个框架（比如PyTorch）练手。把常见的Conv2D、BatchNorm、ReLU这些算子摸透，再扩展到其他框架。别一上来就想全支持，容易把自己搞懵。

好了，关于三个主流框架的模型表示和IR，我就聊这么多。记住一点：不管框架怎么变，算子映射的核心就是IR解析 + 硬件适配。把这两个基本功打扎实了，什么框架来了都不怕。

📚 课程目录

3. 主流AI框架介绍：TensorFlow、PyTorch、MindSpore的模型表示与中间表示（IR）

3.1 TensorFlow：静态图的坚守者

3.2 PyTorch：动态图的灵活派

3.3 MindSpore：全场景的后来者

3.4 三者的IR对比

3.5 从IR到NPU算子的映射流程