1. 初步计划
- 基于预训练模型,在自己的图片上进行推理测试
- 导出 ONNX 格式文件,测试转 MNN 流程,了解 MNN 是如何转换的(待完成)
- **导出 ONNX 格式文件,查看其是如何转换为onnx文件的
-
- 1、通过 API 看 过程
-
-
- 1.1、Pytorch 对 onnx 算子的支持(先跳过)2023-08-15,09:53:59
-
-
- 2、通过 看流程中的 TorchScript 是如何转换的
-
-
- 1.1、Torch jit tracer 实现解析
- 1.2、jit 中的 subgraph rewriter
- 1.3、Torch jit 中的别名分析
-
- 手写脚本转换(未完成)
- 测试转换后的精度(参考转换过程实现工具)
1:LeNet - 5 验证码 训练模型 完成转pt --> onnx
1.1: LeNet - 5该模型的原理
模型原理可见:神经网络 LeNet-5详解
1.1.1:完成一个数据集合训练
1.1.1.1:MINIST训练集合
使用已经存在的训练集合训练:
最终测试 训练的模型
ini
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
import time
from matplotlib import pyplot as plt
import cv2
# 下载并加载数据,并做出一定的预先处理
pipline_train = transforms.Compose([
# 随机旋转图片
transforms.RandomHorizontalFlip(),
# 将图片尺寸resize到32x32
transforms.Resize((32, 32)),
# 将图片转化为Tensor格式
transforms.ToTensor(),
# 正则化(当模型出现过拟合的情况时,用来降低模型的复杂度)
transforms.Normalize((0.1307,), (0.3081,))
])
pipline_test = transforms.Compose([
# 将图片尺寸resize到32x32
transforms.Resize((32, 32)),
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
# 下载数据集
train_set = datasets.MNIST(root="./data", train=True, download=True, transform=pipline_train)
test_set = datasets.MNIST(root="./data", train=False, download=True, transform=pipline_test)
# 加载数据集
train_loadr = torch.utils.data.DataLoader(train_set, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_set, batch_size=32, shuffle=False)
# 搭建 LeNet-5 神经网络结构,并定义前向传播的过程
class LeNet(nn.Module):
def __init__(self):
super(LeNet, self).__init__()
self.conv1 = nn.Conv2d(1, 6, 5)
self.relu = nn.ReLU()
self.max_pool1 = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 16, 5)
self.max_pool2 = nn.MaxPool2d(2, 2)
self.fc1 = nn.Linear(16 * 5 * 5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
x = self.conv1(x)
x = self.relu(x)
x = self.max_pool1(x)
x = self.conv2(x)
x = self.max_pool2(x)
x = x.view(-1, 16 * 5 * 5)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
output = F.log_softmax(x, dim=1)
return output
# 将定义好的网络结构搭载到 GPU/CPU,并定义优化器
# 创建模型,部署gpu
device = torch.device("cpu")
model = LeNet().to(device)
# 定义优化器
optimizer = optim.Adam(model.parameters(), lr=0.001)
# 定义训练过程
def train_runner(model, device, train_loadr, optimizer, epoch):
# 训练模型, 启用 BatchNormalization 和 Dropout, 将BatchNormalization和Dropout置为True
model.train()
total = 0
correct = 0.0
# enumerate迭代已加载的数据集,同时获取数据和数据下标
for i, data in enumerate(train_loadr, 0):
inputs, labels = data
# 把模型部署到device上
inputs, labels = inputs.to(device), labels.to(device)
# 初始化梯度
optimizer.zero_grad()
# 保存训练结果
outputs = model(inputs)
# 计算损失和
# 多分类情况通常使用cross_entropy(交叉熵损失函数), 而对于二分类问题, 通常使用sigmoid
loss = F.cross_entropy(outputs, labels)
# 获取最大概率的预测结果
# dim=1表示返回每一行的最大值对应的列下标
predict = outputs.argmax(dim=1)
total += labels.size(0)
correct += (predict == labels).sum().item()
# 反向传播
loss.backward()
# 更新参数
optimizer.step()
if i % 1000 == 0:
# loss.item()表示当前loss的数值
print(
"Train Epoch{} \t Loss: {:.6f}, accuracy: {:.6f}%".format(epoch, loss.item(), 100 * (correct / total)))
Loss.append(loss.item())
Accuracy.append(correct / total)
return loss.item(), correct / total
def test_runner(model, device, train_loadr):
# 模型验证, 必须要写, 否则只要有输入数据, 即使不训练, 它也会改变权值
# 因为调用eval()将不启用 BatchNormalization 和 Dropout, BatchNormalization和Dropout置为False
model.eval()
# 统计模型正确率, 设置初始值
correct = 0.0
test_loss = 0.0
total = 0
# torch.no_grad将不会计算梯度, 也不会进行反向传播
with torch.no_grad():
for data, label in train_loadr:
data, label = data.to(device), label.to(device)
output = model(data)
test_loss += F.cross_entropy(output, label).item()
predict = output.argmax(dim=1)
# 计算正确数量
total += label.size(0)
correct += (predict == label).sum().item()
# 计算损失值
print("test_avarage_loss: {:.6f}, accuracy: {:.6f}%".format(test_loss / total, 100 * (correct / total)))
# 调用
epoch = 5
Loss = []
Accuracy = []
for epoch in range(1, epoch + 1):
print("start_time", time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time())))
loss, acc = train_runner(model, device, train_loadr, optimizer, epoch)
Loss.append(loss)
Accuracy.append(acc)
test_runner(model, device, train_loadr)
print("end_time: ", time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time())), '\n')
print('Finished Training')
plt.subplot(2, 1, 1)
plt.plot(Loss)
plt.title('Loss')
plt.show()
plt.subplot(2, 1, 2)
plt.plot(Accuracy)
plt.title('Accuracy')
plt.show()
print(model)
torch.save(model, './models/model-mnist.pth') # 保存模型
torch.save(model, './models/model-mnist.pt') # 保存模型
# 导出为onnx格式
batch_size = 1 # 批处理大小
torch_model = torch.load('./models/model-mnist.pth') # 加载模型
# set the model to inference mode
device = torch.device('cpu')
torch_model = model.to(device)
torch_model.eval()
x = torch.randn(batch_size, 1, 32, 32) # 生成张量
export_onnx_file = "./models/model-mnist.onnx" # 目的ONNX文件名
torch.onnx.export(torch_model,
x,
export_onnx_file,
opset_version=10,
do_constant_folding=True, # 是否执行常量折叠优化
input_names=["input"], # 输入名
output_names=["output"], # 输出名
dynamic_axes={"input": {0: "batch_size"},
"output": {0: "batch_size"}})
if __name__ == '__main__':
device = torch.device('cpu')
model = torch.load('./models/model-mnist.pth') # 加载模型
model = model.to(device)
model.eval() # 把模型转为test模式
# 获取所有参数张量(权重和偏置)
all_parameters = list(model.parameters())
# 获取权重张量
weight_tensors = [param for param in all_parameters if len(param.shape) > 1]
# 打印权重张量的形状
for idx, weight_tensor in enumerate(weight_tensors):
print(f"打印权重张量的形状: Weight Tensor {idx + 1} shape: {weight_tensor.shape}")
print(f'获取所有参数张量(权重和偏置): {weight_tensors}')
# # 读取要预测的图片
# img = cv2.imread("E:\LeNet-5_For_test\data\MNIST\raw\test\7\1141.png")
# img = cv2.resize(img, dsize=(32, 32), interpolation=cv2.INTER_NEAREST)
# 读取要预测的图片
img = cv2.imread("E:\LeNet-5_For_test\picture\captcha_images_v2\captcha_images_v2\2bg48.png")
img = cv2.resize(img, dsize=(32, 32), interpolation=cv2.INTER_NEAREST)
plt.imshow(img, cmap="gray") # 显示图片
plt.axis('off') # 不显示坐标轴
plt.show()
# 导入图片,图片扩展后为[1,1,32,32]
trans = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) # 图片转为灰度图,因为mnist数据集都是灰度图
img = trans(img)
img = img.to(device)
img = img.unsqueeze(0) # 图片扩展多一维,因为输入到保存的模型中是4维的[batch_size,通道,长,宽],而普通图片只有三维,[通道,长,宽]
# 预测
output = model(img)
prob = F.softmax(output, dim=1) # prob是10个分类的概率
print("概率:", prob)
value, predicted = torch.max(output.data, 1)
predict = output.argmax(dim=1)
print("预测类别:", predict.item())



1.1.1.2:验证码手动下载训练(待完成,文件下载完成,标注未完成)
2:训练一个识别验证码的集合(未完成)
3:手动实现一个转换(了解原理中见下文)
如何从.pt 转为 .onnx 的?
shell
# 导出为onnx格式
# batch_size = 1 # 批处理大小
torch_model = torch.load('./models/model-mnist.pth') # 加载模型
# set the model to inference mode
# device = torch.device('cpu')
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# torch_model = model.to(device)
# torch_model.eval()
#
# input_dumpy = torch.randn(batch_size, 1, 32, 32) # 生成张量
# export_onnx_file = "./models/model-mnist.onnx" # 目的ONNX文件名
# torch.onnx.export(torch_model,
# input_dumpy,
# export_onnx_file,
# opset_version=10,
# do_constant_folding=True, # 是否执行常量折叠优化
# input_names=["input"], # 输入名
# output_names=["output"], # 输出名
# dynamic_axes={"input": {0: "batch_size"},
# "output": {0: "batch_size"}})
1、了解如何部署模型
zhuanlan.zhihu.com/p/477743341
2、部署模型的问题
最终的一个生成的中间件onnx如下所示:

可以和代码对比的看到是一致的

每个算子记录了算子属性、图结构、权重三类信息。
- 算子属性信息即图中 attributes 里的信息,对于卷积来说,算子属性包括了卷积核大小(kernel_shape)、卷积步长(strides)等内容。这些算子属性最终会用来生成一个具体的算子。
- 图结构信息指算子节点在计算图中的名称、邻边的信息。对于图中的卷积来说,该算子节点叫做 Conv_2,输入数据叫做 11,输出数据叫做 12。根据每个算子节点的图结构信息,就能完整地复原出网络的计算图。
- 权重信息指的是网络经过训练后,算子存储的权重信息。对于卷积来说,权重信息包括卷积核的权重值和卷积后的偏差值。点击图中 conv1.weight, conv1.bias 后面的加号即可看到权重信息的具体内容。
现在,我们有了 ONNX 模型。让我们看看最后该如何把这个模型运行起来。
运行之后如果想要修改参数,例如图像的大小只能从 刚开始设计的时候写入,在后续中如果要求修改的化是没有办法的,这里就需要适配算法和算子。
如何具体的适配算法和算子需要具体的查看是那部分的内容。
模型部署中常见的几类困难有:模型的动态化;新算子的实现;框架间的兼容。
pt 是如何转为 onnx
详细的解析 torch.onnx.export
实现原理来更好地应对该函数的报错(由于模型部署的兼容性问题,部署复杂模型时该函数时常会报错)。计算图导出方法。
TorchScript 是一种序列化和优化 PyTorch 模型的格式,在优化过程中,一个torch.nn.Module模型会被转换成 TorchScript 的 torch.jit.ScriptModule模型。现在, TorchScript 也被常当成一种中间表示使用。其他文章中对 TorchScript 有详细的介绍(zhuanlan.zhihu.com/p/486914187)

了解导出模型的两种方法
跟踪法只能通过实际运行一遍模型的方法导出模型的静态图,即无法识别出模型中的控制流(如循环);
记录法则能通过解析模型来正确记录所有的控制流。
使用代码来导出两份不同的过程
期望结果是导出两份.onnx 文件但是在执行时出现了报错,查看:
这个建议是关于在 PyTorch 中使用模型转换技术(tracing 和 scripting)的使用情况。
在 PyTorch 中,模型转换是将动态图模型转换为静态图模型的过程,以便在不需要 Python 解释器的情况下运行。其中有两种主要的转换方式:tracing 和 scripting。
- Tracing :使用 torch.jit.trace 方法,你可以通过运行模型的一个示例输入,记录其计算图。然后,这个计算图可以被优化和导出到其他框架(如 ONNX)中。但是,trace 是基于一个示例输入的,如果模型的行为在不同输入上有所不同,可能会导致跟踪的计算图不准确。
- Scripting :使用 torch.jit.script 方法,你可以将整个模型的代码转换为静态图。这允许模型在不同输入上都使用相同的计算图,从而避免了在不同输入上产生不准确的计算图。
在你提到的建议中,它指出对于某些情况,tracing 并不适合,而使用 scripting 更合适。如果你使用 trace 来记录这些模型,可能会在以后的模型调用中悄无声息地得到不正确的结果。当进行可能导致不正确的跟踪时,跟踪器会尝试发出警告。
这个建议的要点是:对于模型的行为依赖于不同输入的情况,使用 scripting 更可靠,因为它会将整个模型的代码转换为静态图,而不仅仅是基于一个示例输入的计算图。这可以确保模型在不同输入上都有一致的行为,避免潜在的错误。
Q : 跟踪法和记录法在导出带控制语句的计算图时有什么区别?
A:"跟踪法"和"记录法"在导出带有控制语句(如循环和条件语句)的计算图时有一些区别。这两种方法在这种情况下的差异:
- 跟踪法(Tracing) :
跟踪法是通过执行一个模型的示例输入来记录计算图。在有控制语句的情况下,跟踪法会记录在示例输入上执行的控制流路径。然而,跟踪法可能只会记录示例输入的某个特定路径,导致只有特定的控制流路径被捕获。
如果模型的控制流在不同的输入上有所不同,那么通过跟踪法记录的计算图可能会在其他输入上产生错误的结果。这是因为只有一个示例路径被记录下来,而其他路径可能会被忽略。
- 记录法(Scripting) :
记录法是通过将整个模型的代码转换为静态图来记录计算图。在有控制语句的情况下,记录法会将整个控制流的逻辑转换为静态图的形式。这意味着所有可能的控制流路径都会被捕获,不会因为输入而有所不同。
记录法能够处理更复杂的模型逻辑,尤其是涉及多个分支、循环等控制结构的情况。因为整个模型的代码被转换为静态图,所以不会因为输入的不同而产生不一致的行为。
总之,当模型涉及复杂的控制流结构时,使用"记录法"(torch.jit.script )会更加可靠,因为它可以捕获所有的控制流路径,避免了在不同输入上产生不准确的计算图。而"跟踪法"(torch.jit.trace)可能只捕获特定路径的计算图,因此在涉及不同输入情况下可能会出现问题。
PyTorch 对 ONNX 的算子支持(先跳过,后续再看先搞懂如何转换的)
在转换普通的torch.nn.Module模型时,PyTorch 一方面会用跟踪法执行前向推理,把遇到的算子整合成计算图;另一方面,PyTorch 还会把遇到的每个算子翻译成 ONNX 中定义的算子。在这个翻译过程中,可能会碰到以下情况:
- 该算子可以一对一地翻译成一个 ONNX 算子。
- 该算子在 ONNX 中没有直接对应的算子,会翻译成一至多个 ONNX 算子。
- 该算子没有定义翻译成 ONNX 的规则,报错。
那么,该如何查看 PyTorch 算子与 ONNX 算子的对应情况呢?由于 PyTorch 算子是向 ONNX 对齐的,这里我们先看一下 ONNX 算子的定义情况,再看一下 PyTorch 定义的算子映射关系。
ONNX 算子文档
ONNX 算子的定义情况,都可以在官方的算子文档中查看
如何自定义算子实现
PyTorch 算子顺利转换到 ONNX ,我们需要保证以下三个环节都不出错:
- 算子在 PyTorch 中有实现
- 有把该 PyTorch 算子映射成一个或多个 ONNX 算子的方法
- ONNX 有相应的算子
可在实际部署中,这三部分的内容都可能有所缺失。其中最坏的情况是:我们定义了一个全新的算子,它不仅缺少 PyTorch 实现,还缺少 PyTorch 到 ONNX 的映射关系。对于这三个环节,我们也分别都有以下的添加支持的方法:
- PyTorch 算子
-
- 组合现有算子
- 添加 TorchScript 算子
- 添加普通 C++ 拓展算子
- 映射方法
-
- 为 ATen 算子添加符号函数
- 为 TorchScript 算子添加符号函数
- 封装成 torch.autograd.Function 并添加符号函数
- ONNX 算子
-
- 使用现有 ONNX 算子
- 定义新 ONNX 算子
那么面对不同的情况时,就需要我们灵活地选用和组合这些方法。听起来是不是很复杂?别担心,本篇文章中,我们将围绕着三种算子映射方法,学习三个添加算子支持的实例,来理清如何合适地为 PyTorch 算子转 ONNX 算子的三个环节添加支持。
算子支持可以详细查看官方文档:pytorch.org/cppdocs/#at...
pytorch c++ API 主要 解决一下几个问题:
- ATen:基础张量和数学运算库,其他所有内容都建立在其上。也就是说pytorch中支持的基础算子
- Autograd:通过自动微分增强 ATen。
- C++ 前端:用于训练和评估机器学习模型的高级构造。
- TorchScript:TorchScript JIT 编译器和解释器的接口。
- C++ 扩展:一种使用自定义 C++ 和 CUDA 例程扩展 Python API 的方法。
了解 TorchScript
官方动态图解释了
矩阵乘法,然后做加法,在做双曲正切函数作为损失函数,然后计算loss梯度

TorchScript 就是为了解决这个问题而诞生的工具。包括代码的追踪及解析、中间表示的生成、模型优化、序列化等各种功能,可以说是覆盖了模型部署的方方面面。
模型转换
作为模型部署的一个范式,通常我们都需要生成一个模型的中间表示(IR),这个 IR 拥有相对固定的图结构,所以更容易优化
以我们之前的代码中添加以下内容:
python
# 通过trace的方法生成IR需要一个输入样例
dummy_input = torch.rand(1, 1, 32, 32)
# IR生成
with torch.no_grad():
jit_model = torch.jit.trace(model, dummy_input)
print('jit_model 如下所示: ')
print(jit_model)
print('jit_model.graph 如下:')
print(jit_model.graph)
print('jit_model.code 如下所示:')
print(jit_model.code)
查看生成的日志可见:
ini
jit_model 如下所示:
LeNet(
original_name=LeNet
(conv1): Conv2d(original_name=Conv2d)
(relu): ReLU(original_name=ReLU)
(max_pool1): MaxPool2d(original_name=MaxPool2d)
(conv2): Conv2d(original_name=Conv2d)
(max_pool2): MaxPool2d(original_name=MaxPool2d)
(fc1): Linear(original_name=Linear)
(fc2): Linear(original_name=Linear)
(fc3): Linear(original_name=Linear)
)
jit_model.graph 如下:
graph(%self.1 : __torch__.LeNet,
%x.1 : Float(1, 1, 32, 32, strides=[1024, 1024, 32, 1], requires_grad=0, device=cpu)):
%fc3 : __torch__.torch.nn.modules.linear.___torch_mangle_3.Linear = prim::GetAttr[name="fc3"](%self.1)
%fc2 : __torch__.torch.nn.modules.linear.___torch_mangle_2.Linear = prim::GetAttr[name="fc2"](%self.1)
%fc1 : __torch__.torch.nn.modules.linear.Linear = prim::GetAttr[name="fc1"](%self.1)
%max_pool2 : __torch__.torch.nn.modules.pooling.___torch_mangle_1.MaxPool2d = prim::GetAttr[name="max_pool2"](%self.1)
%conv2 : __torch__.torch.nn.modules.conv.___torch_mangle_0.Conv2d = prim::GetAttr[name="conv2"](%self.1)
%max_pool1 : __torch__.torch.nn.modules.pooling.MaxPool2d = prim::GetAttr[name="max_pool1"](%self.1)
%relu : __torch__.torch.nn.modules.activation.ReLU = prim::GetAttr[name="relu"](%self.1)
%conv1 : __torch__.torch.nn.modules.conv.Conv2d = prim::GetAttr[name="conv1"](%self.1)
%151 : Tensor = prim::CallMethod[name="forward"](%conv1, %x.1)
%152 : Tensor = prim::CallMethod[name="forward"](%relu, %151)
%153 : Tensor = prim::CallMethod[name="forward"](%max_pool1, %152)
%154 : Tensor = prim::CallMethod[name="forward"](%conv2, %153)
%155 : Tensor = prim::CallMethod[name="forward"](%max_pool2, %154)
%105 : int = prim::Constant[value=-1]() # E:\LeNet-5_For_test\torchScrpitStudy.py:56:0
%106 : int = prim::Constant[value=400]() # E:\LeNet-5_For_test\torchScrpitStudy.py:56:0
%107 : int[] = prim::ListConstruct(%105, %106)
%input.9 : Float(1, 400, strides=[400, 1], requires_grad=0, device=cpu) = aten::view(%155, %107) # E:\LeNet-5_For_test\torchScrpitStudy.py:56:0
%156 : Tensor = prim::CallMethod[name="forward"](%fc1, %input.9)
%input.13 : Float(1, 120, strides=[120, 1], requires_grad=0, device=cpu) = aten::relu(%156) # C:\Users\user\AppData\Local\Programs\Python\Python311\Lib\site-packages\torch\nn\functional.py:1457:0
%157 : Tensor = prim::CallMethod[name="forward"](%fc2, %input.13)
%input.17 : Float(1, 84, strides=[84, 1], requires_grad=0, device=cpu) = aten::relu(%157) # C:\Users\user\AppData\Local\Programs\Python\Python311\Lib\site-packages\torch\nn\functional.py:1457:0
%158 : Tensor = prim::CallMethod[name="forward"](%fc3, %input.17)
%114 : int = prim::Constant[value=1]() # C:\Users\user\AppData\Local\Programs\Python\Python311\Lib\site-packages\torch\nn\functional.py:1932:0
%115 : NoneType = prim::Constant()
%116 : Float(1, 10, strides=[10, 1], requires_grad=0, device=cpu) = aten::log_softmax(%158, %114, %115) # C:\Users\user\AppData\Local\Programs\Python\Python311\Lib\site-packages\torch\nn\functional.py:1932:0
return (%116)
jit_model.code 如下所示:
def forward(self,
x: Tensor) -> Tensor:
fc3 = self.fc3
fc2 = self.fc2
fc1 = self.fc1
max_pool2 = self.max_pool2
conv2 = self.conv2
max_pool1 = self.max_pool1
relu = self.relu
conv1 = self.conv1
_0 = (relu).forward((conv1).forward(x, ), )
_1 = (conv2).forward((max_pool1).forward(_0, ), )
input = torch.view((max_pool2).forward(_1, ), [-1, 400])
input0 = torch.relu((fc1).forward(input, ))
input1 = torch.relu((fc2).forward(input0, ))
_2 = torch.log_softmax((fc3).forward(input1, ), 1)
return _2
主要做几点,
1、生成一个模型的网络层
2、生成该算法的一个 graph TorchScript 有它自己对于 Graph 以及其中元素的定义,目前来讲看不太懂。
3、TorchScript 的 IR 是可以还原成 python 代码的,如果你生成了一个 TorchScript 模型并且想知道它的内容对不对,那么可以通过这样的方式来做一些简单的检查。
模型优化
上面的可视化中只有 LeNet-5 里 forward 的部分,而且顺序似乎还是不太一样的。那么其中的子模块信息是不是丢失了呢?如果没有丢失,那么怎么样才能确定子模块的内容是否正确呢?
TorchScript 支持对网络的优化吗,这里我们就可以用一个pass解决这个问题:
scss
print('jit_model.code 如下所示:')
print(jit_model.code)
# 调用inline pass,对graph做变换
print("调用inline pass,对graph做变换")
torch._C._jit_pass_inline(jit_model.graph)
print(jit_model.code)
查看打印的日志对照发现
ini
jit_model.code 如下所示:
def forward(self,
x: Tensor) -> Tensor:
fc3 = self.fc3
fc2 = self.fc2
fc1 = self.fc1
max_pool2 = self.max_pool2
conv2 = self.conv2
max_pool1 = self.max_pool1
relu = self.relu
conv1 = self.conv1
_0 = (relu).forward((conv1).forward(x, ), )
_1 = (conv2).forward((max_pool1).forward(_0, ), )
input = torch.view((max_pool2).forward(_1, ), [-1, 400])
input0 = torch.relu((fc1).forward(input, ))
input1 = torch.relu((fc2).forward(input0, ))
_2 = torch.log_softmax((fc3).forward(input1, ), 1)
return _2
调用inline pass,对graph做变换
def forward(self,
x: Tensor) -> Tensor:
fc3 = self.fc3
fc2 = self.fc2
fc1 = self.fc1
max_pool2 = self.max_pool2
conv2 = self.conv2
max_pool1 = self.max_pool1
relu = self.relu
conv1 = self.conv1
bias = conv1.bias
weight = conv1.weight
input = torch._convolution(x, weight, bias, [1, 1], [0, 0], [1, 1], False, [0, 0], 1, False, False, True, True)
input0 = torch.relu(input)
input1 = torch.max_pool2d(input0, [2, 2], [2, 2], [0, 0], [1, 1])
bias0 = conv2.bias
weight0 = conv2.weight
input2 = torch._convolution(input1, weight0, bias0, [1, 1], [0, 0], [1, 1], False, [0, 0], 1, False, False, True, True)
x0 = torch.max_pool2d(input2, [2, 2], [2, 2], [0, 0], [1, 1])
input3 = torch.view(x0, [-1, 400])
bias1 = fc1.bias
weight1 = fc1.weight
input4 = torch.linear(input3, weight1, bias1)
input5 = torch.relu(input4)
bias2 = fc2.bias
weight2 = fc2.weight
input6 = torch.linear(input5, weight2, bias2)
input7 = torch.relu(input6)
bias3 = fc3.bias
weight3 = fc3.weight
input8 = torch.linear(input7, weight3, bias3)
return torch.log_softmax(input8, 1)
上面代码中我们使用了一个名为inline的pass,将所有子模块进行内联,这样我们就能看见更完整的推理代码。pass是一个来源于编译原理的概念,一个 TorchScript 的 pass 会接收一个图,遍历图中所有元素进行某种变换,生成一个新的图。我们这里用到的inline起到的作用就是将模块调用展开,尽管这样做并不能直接影响执行效率,但是它其实是很多其他pass的基础。PyTorch 中定义了非常多的 pass 来解决各种优化任务。
在编译原理中,"pass" 通常指的是编译器的一个阶段,该阶段对源代码进行转换和优化,以便生成更高效、更优化的目标代码。每个 "pass" 都是一个特定的转换或优化过程,它在编译过程的特定阶段应用于代码,并在多个 "pass" 的迭代中逐步改进代码的质量和性能。
每个 "pass" 可以执行以下操作之一或多个:
- 分析(Analysis) :在这个阶段,编译器收集关于源代码的信息,例如符号表、数据流分析、控制流分析等。这些信息有助于后续的优化过程。
- 转换(Transformation) :这个阶段会对源代码进行修改,以便在保持语义不变的情况下改进代码的性能、可读性或其他方面。例如,常数折叠、循环展开、内联等。
- 优化(Optimization) :在此阶段,编译器尝试改进代码的执行效率,以减少运行时开销。优化可以涉及复杂的技术,如循环优化、数据局部性优化、寄存器分配等。
- 代码生成(Code Generation) :在此阶段,编译器根据经过分析和优化的中间表示生成目标代码,可以是汇编语言、机器码等。
编译器通常会应用多个 "pass",并且它们的顺序和数量可能会根据编译器的设计和目标语言的特性而变化。通过不同 "pass" 的迭代,编译器可以逐步改进源代码的性能和质量。
序列化
不管是哪种方法创建的 TorchScript 都可以进行序列化,比如:
ini
# 将模型序列化
jit_model.save('jit_model.pth')
# 加载序列化后的模型
jit_model = torch.jit.load('jit_model.pth')
序列化后的模型不再与 python 相关,可以被部署到各种平台上。
PyTorch 提供了可以用于 TorchScript 模型推理的 c++ API,序列化后的模型终于可以不依赖 python 进行推理了:
arduino
// 加载生成的torchscript模型
auto module = torch::jit::load('jit_model.pth');
// 根据任务需求读取数据
std::vector<torch::jit::IValue> inputs = ...;
// 计算推理结果
auto output = module.forward(inputs).toTensor();
TorchScript总结:
TorchScript 的主要用途是进行模型部署,需要记录生成一个便于推理优化的 IR,对计算图的编辑通常都是面向性能提升等等,不会给模型本身添加新的功能。
python
import torch.jit
class CustomModel(torch.jit.ScriptModule):
@torch.jit.script_method
def forward(self, x):
x = x * 2
x.add_(0)
# 输入张量 x 改变形状为一维向量
x = x.view(-1)
output_condition = self.condition(x)
return output_condition
@torch.jit.script_method
def condition(self, x):
if x[0] > 1:
return x[0]
else:
return x[-1]
def __getstate__(self):
return None
# 实例化自定义模型
model = CustomModel()
# 创建一个示例输入
input_example = torch.tensor([1.0, 2.0, 3.0]) # 示例输入
# 在模型上进行前向传播
output = model(input_example)
# save model data
torch.save(model, 'test_for_function_model.pt')
torch.onnx.export(model, input_example, "test_for_function_model.onnx", verbose=True)
ini
Exported graph: graph(%x.1 : Float(3, strides=[1], requires_grad=0, device=cpu)):
%/Constant_output_0 : Long(device=cpu) = onnx::Constant[value={0}, onnx_name="/Constant"](), scope: CustomModel::
%/Constant_1_output_0 : Long(device=cpu) = onnx::Constant[value={-1}, onnx_name="/Constant_1"](), scope: CustomModel::
%/Constant_2_output_0 : Float(requires_grad=0, device=cpu) = onnx::Constant[value={2}, onnx_name="/Constant_2"](), scope: CustomModel:: # E:\LeNet-5_For_test\testForForwardFunction.py:7:12
%/Mul_output_0 : Float(3, strides=[1], device=cpu) = onnx::Mul[onnx_name="/Mul"](%x.1, %/Constant_2_output_0), scope: CustomModel:: # E:\LeNet-5_For_test\testForForwardFunction.py:7:12
%/Constant_3_output_0 : Float(requires_grad=0, device=cpu) = onnx::Constant[value={0}, onnx_name="/Constant_3"](), scope: CustomModel:: # E:\LeNet-5_For_test\testForForwardFunction.py:8:8
%/Add_output_0 : Float(3, strides=[1], device=cpu) = onnx::Add[onnx_name="/Add"](%/Mul_output_0, %/Constant_3_output_0), scope: CustomModel:: # E:\LeNet-5_For_test\testForForwardFunction.py:8:8
%/Constant_4_output_0 : Long(1, strides=[1], device=cpu) = onnx::Constant[value={-1}, onnx_name="/Constant_4"](), scope: CustomModel:: # E:\LeNet-5_For_test\testForForwardFunction.py:9:12
%/Reshape_output_0 : Float(3, strides=[1], device=cpu) = onnx::Reshape[allowzero=0, onnx_name="/Reshape"](%/Add_output_0, %/Constant_4_output_0), scope: CustomModel:: # E:\LeNet-5_For_test\testForForwardFunction.py:9:12
%/Gather_output_0 : Float(device=cpu) = onnx::Gather[axis=0, onnx_name="/Gather"](%/Reshape_output_0, %/Constant_output_0), scope: CustomModel:: # E:\LeNet-5_For_test\testForForwardFunction.py:15:11
%/Constant_5_output_0 : Float(requires_grad=0, device=cpu) = onnx::Constant[value={1}, onnx_name="/Constant_5"](), scope: CustomModel:: # E:\LeNet-5_For_test\testForForwardFunction.py:15:11
%/Greater_output_0 : Bool(device=cpu) = onnx::Greater[onnx_name="/Greater"](%/Gather_output_0, %/Constant_5_output_0), scope: CustomModel:: # E:\LeNet-5_For_test\testForForwardFunction.py:15:11
%/Cast_output_0 : Bool(device=cpu) = onnx::Cast[to=9, onnx_name="/Cast"](%/Greater_output_0), scope: CustomModel:: # E:\LeNet-5_For_test\testForForwardFunction.py:15:8
%output_condition : Float(device=cpu) = onnx::If[onnx_name="/If"](%/Cast_output_0), scope: CustomModel:: # E:\LeNet-5_For_test\testForForwardFunction.py:15:8
block0():
%/Gather_1_output_0 : Float(device=cpu) = onnx::Gather[axis=0, onnx_name="/Gather_1"](%/Reshape_output_0, %/Constant_output_0), scope: CustomModel:: # E:\LeNet-5_For_test\testForForwardFunction.py:16:19
-> (%/Gather_1_output_0)
block1():
%/Gather_2_output_0 : Float(device=cpu) = onnx::Gather[axis=0, onnx_name="/Gather_2"](%/Reshape_output_0, %/Constant_1_output_0), scope: CustomModel:: # E:\LeNet-5_For_test\testForForwardFunction.py:18:19
-> (%/Gather_2_output_0)
return (%output_condition)
============= Diagnostic Run torch.onnx.export version 2.0.1+cu118 =============
verbose: False, log level: Level.ERROR
======================= 0 NONE 0 NOTE 0 WARNING 0 ERROR ========================
生成的图片如下所示:

了解文件结构如上:
Jit trace 在 python 侧的接口为torch.jit.trace,输入的参数会经过层层传递,最终会进入torch/jit/frontend/trace.cpp中的trace函数中。
scss
std::pair<std::shared_ptr<TracingState>, Stack> trace(
Stack inputs,
const std::function<Stack(Stack)>& traced_fn,
std::function<std::string(const Variable&)> var_name_lookup_fn,
bool strict,
bool force_outplace,
Module* self,
const std::vector<std::string>& argument_names) {
try {
// Start tracing, treating 'inputs' as inputs to the trace, which can be
// varied on subsequent invocations of the trace. Any other variables
// will be treated as constants.
if (isTracing()) {
AT_ERROR("Tracing can't be nested");
}
auto state = std::make_shared<TracingState>();
setTracingState(state);
// if we are a module, then make sure the modules parameters are in the map
// and mapped to accesses to the self object
if (self) {
Value* self_value = state->graph->insertInput(0, "self")->setType(
self->_ivalue()->type());
gatherParametersAndBuffers(state, self_value, *self, {"__module"});
}
// When enough argument name hints are provided, use them as debug names
// for traced function/modules.
// Here argument_names is allowed to have more names than needed because
// some arguments may have valid default values, therefore they don't need
// example inputs.
if (argument_names.size() >= inputs.size()) {
for (size_t i = 0, e = inputs.size(); i < e; ++i) {
IValue& input = inputs[i];
input = addInput(
state,
input,
input.type(),
state->graph->addInput(argument_names[i]));
}
} else {
for (IValue& input : inputs) {
input = addInput(state, input, input.type(), state->graph->addInput());
}
}
auto graph = state->graph;
getTracingState()->lookup_var_name_fn = std::move(var_name_lookup_fn);
getTracingState()->strict = strict;
getTracingState()->force_outplace = force_outplace;
// Invoke the traced function
auto out_stack = traced_fn(inputs);
// Exit a trace, treating 'out_stack' as the outputs of the trace. These
// are the variables whose values will be computed upon subsequent
// invocations of the trace.
size_t i = 0;
for (auto& output : out_stack) {
// NB: The stack is in "reverse" order, so when we pass the diagnostic
// number we need to flip it based on size.
state->graph->registerOutput(
state->getOutput(output, out_stack.size() - i));
i++;
}
setTracingState(nullptr);
if (getInlineEverythingMode()) {
Inline(*graph);
}
FixupTraceScopeBlocks(graph, self);
NormalizeOps(graph);
return {state, out_stack};
} catch (...) {
tracer::abandon();
throw;
}
}
上面给出的是 trace 函数 在c++中对应的源码,里面有很多部分缺失,做了封装,具体流程可以参见该博客:
zhuanlan.zhihu.com/p/489090393针对trace 机制的解读
ONNX 模型在底层是用什么格式存储的
二进制方式Protobuf ,
- ModelProto
-
- GraphProto
-
-
- NodeProto
- ValueInfoProto
-
ONNX导出时执行那些函数
参考该博客:zhuanlan.zhihu.com/p/489090393
- 加载 ops 的 symbolic 函数,主要是 torch 中预定义的 symbolic。主要做一个映射
- 设置环境,包括 opset_version,是否折叠常量等等。
- 使用 jit trace 生成 Graph。(很复杂可以查看c++代码和博客)
- 将 Graph 中的 Node 映射成 ONNX 的 Node,并进行必要的优化。
- 将模型导出成 ONNX 的序列化格式。
对比精度

使用Debugger方法来插桩打点确认精度
由于在 转换是无法判断哪一层会出现精度丢失,所以可以根据ONNX算子的特性
定义一个叫做 Debug 的 ONNX 算子,它有一个属性调试名 name。而由于每一个 ONNX 算子节点又自带了输出张量的名称,这样一来,ONNX 节点的输出名和调试名绑定在了一起。我们可以顺着 PyTorch 里的调试名,找到对应 ONNX 里的输出,完成 PyTorch 和 ONNX 的对应。

算子的添加可以参考这个文章:
zhuanlan.zhihu.com/p/513387413
主要介绍的 API为 g.op()
符号函数,可以看成是 PyTorch 算子类的一个静态方法。在把 PyTorch 模型转换成 ONNX 模型时,各个 PyTorch 算子的符号函数会被依次调用,以完成 PyTorch 算子到 ONNX 算子的转换。符号函数的定义一般如下:
less
def symbolic(g: torch._C.Graph, input_0: torch._C.Value, input_1: torch._C.Value, ...):
其中,torch._C.Graph 和 torch._C.Value 都对应 PyTorch 的 C++ 实现里的一些类。只需要知道第一个参数就固定叫 g,它表示和计算图相关的内容;后面的每个参数都表示算子的输入,需要和算子的前向推理接口的输入相同。对于 ATen 算子来说,它们的前向推理接口就是上述两个 .pyi 文件里的函数接口。
g 有一个方法 op。在把 PyTorch 算子转换成 ONNX 算子时,需要在符号函数中调用此方法来为最终的计算图添加一个 ONNX 算子。其定义如下:
less
def op(name: str, input_0: torch._C.Value, input_1: torch._C.Value, ...)
其中,第一个参数是算子名称。如果该算子是普通的 ONNX 算子,只需要把它在 ONNX 官方文档里的名称填进去即可。
在最简单的情况下,我们只要把 PyTorch 算子的输入用g.op()一一对应到 ONNX 算子上即可,并把g.op()的返回值作为符号函数的返回值。在情况更复杂时,我们转换一个 PyTorch 算子可能要新建若干个 ONNX 算子。

Debugger 类有三个成员变量:
- torch_value 记录了运行 PyTorch 模型后每个调试张量的值。
- onnx_value 记录了运行 ONNX 模型后每个调试张量的值。
- output_debug_name 记录了把调试张量加入 ONNX 的输出后,每个输出张量的调试名。
稍后我们会在类实现的代码里看到这些成员变量的具体用法。
Debugger 类有以下方法:
- debug 封装了之前编写好的 debug_apply。该方法需要在原 PyTorch 模型中调用,可以为导出的 ONNX 模型添加 Debug 算子节点,同时记录 PyTorch 调试张量值。
- extract_debug_model 和 ONNX 的子模型提取函数的用法类似,可以把带调试节点的 ONNX 模型转化成一个可以输出调试张量的 ONNX 模型。
- run_debug_model 会使用 ONNX Runtime 运行模型,得到 ONNX 调试张量值。
- print_debug_result 会比较 PyTorch 和 ONNX 的调试张量值,输出比较的结果。
精度对齐的核心:
进行推理对比: 使用转换后的 ONNX 模型和原始 PyTorch 模型进行推理,然后比较输出结果。可以计算输出之间的均方误差(MSE)或其他精度指标来评估精度的一致性
代码封装成一个类:debugForModelCompare
python
import torch
import onnx
import onnxruntime
import numpy as np
class DebugOp(torch.autograd.Function):
@staticmethod
def forward(ctx, tensor_x, name):
return tensor_x
@staticmethod
def symbolic(g, tensor_x, name):
return g.op("my::Debug", tensor_x, name_s=name)
# torch.autograd.Function
# PyTorch 的一个可导函数,只要为其定义了前向推理和反向传播的实现,
# 我们就可以把它当成一个普通 PyTorch 函数来使用。
# PyTorch 会自动调度该函数,合适地执行前向和反向计算。
# 对模型部署来说,Function 类有一个很好的性质:
# 如果它定义了 symbolic 静态方法,
# 该 Function 在执行 torch.onnx.export() 时就可以根据 symbolic 中定义的规则转换成 ONNX 算子。
# 这个 symbolic 就是前面提到的符号函数,只是它的名称必须是 symbolic 而已。
#
# 我们在使用 Function 的派生类做推理时,
# 不应该显式地调用 forward(),而应该调用其 apply 方法。
debug_apply = DebugOp.apply
# 原因:
# 因为 .apply() 方法会处理计算图的构建、
# 内存管理等细节,确保一致的行为。
# 而 .forward() 方法是设计用来在 Autograd 过程中被调用的
class Debugger():
def __init__(self):
super().__init__()
self.torch_value = dict()
self.onnx_value = dict()
self.output_debug_name = []
def debug(self, x, name):
# self.torch_value[name] = x.detach().cpu().numpy()
self.torch_value[name] = x.detach().cpu().numpy()
return debug_apply(x, name)
def extract_debug_model(self, input_path, output_path):
model = onnx.load(input_path)
inputs = [input.name for input in model.graph.input]
outputs = []
for node in model.graph.node:
if node.op_type == 'Debug':
# 细节一:onnx是 二进制保存的方式需要转str要编码格式
debug_name = node.attribute[0].s.decode('ASCII')
self.output_debug_name.append(debug_name)
output_name = node.output[0]
outputs.append(output_name)
# 消除onnx不支持自定义Debug节点,定义为Identity
node.op_type = 'Identity'
# 为空会被作为原生节点
node.domain = ''
# 删除该节点的 参数,因为变成了identity节点
del node.attribute[:]
e = onnx.utils.Extractor(model)
extracted = e.extract_model(inputs, outputs)
onnx.save(extracted, output_path)
# 调试张量输出会记录在 debugger.onnx_value
def run_debug_model(self, input, debug_model):
sess = onnxruntime.InferenceSession(debug_model,
providers=['CPUExecutionProvider'])
onnx_outputs = sess.run(None, input)
for name, value in zip(self.output_debug_name, onnx_outputs):
self.onnx_value[name] = value
def print_debug_result(self):
for name in self.torch_value.keys():
if name in self.onnx_value:
squared_diff = (self.torch_value[name] - self.onnx_value[name]) ** 2
mse = np.mean(squared_diff)
print(f"{name} MSE: {mse}")
调用时如下:
ini
debugger = Debugger()
def new_forward(self, x):
x = self.conv1(x)
x = debugger.debug(x, 'x_conv1')
x = self.relu(x)
x = debugger.debug(x, 'x_relu')
x = self.max_pool1(x)
x = debugger.debug(x, 'x_max_pool1')
x = self.conv2(x)
x = debugger.debug(x, 'x_conv2')
x = self.max_pool2(x)
x = debugger.debug(x, 'x_max_pool2')
x = x.view(-1, 16 * 5 * 5)
x = debugger.debug(x, 'x.view(-1, 16 * 5 * 5)')
x = F.relu(self.fc1(x))
x = debugger.debug(x, 'x_F.relu')
x = F.relu(self.fc2(x))
x = debugger.debug(x, 'x_F.relu')
x = self.fc3(x)
x = debugger.debug(x, 'x_self.fc3')
output_forward = F.log_softmax(x, dim=1)
output_forward = debugger.debug(output_forward, 'x_F.log_softmax')
return output_forward
# 替换掉 torch_model 对象的 forward 方法,类似于hook 和 反射 但又不一样
torch_model.forward = MethodType(new_forward, torch_model)
#
dummy_input = torch.randn(1, 1, 32, 32)
torch.onnx.export(torch_model, dummy_input, './debug_models/before_debug.onnx', input_names=['input'])
debugger.extract_debug_model('./debug_models/before_debug.onnx', './debug_models/after_debug.onnx')
debugger.run_debug_model({'input': dummy_input.numpy()}, './debug_models/after_debug.onnx')
print('print debug result')
debugger.print_debug_result()
生成的文件如下:对照原来的LeNet-5 发现在每层插入了debug节点

这些在cpu上训练的结果:
yaml
print debug result
x_conv1 MSE: 3.1806478922137083e-15
x_relu MSE: 1.792849417916239e-15
x_max_pool1 MSE: 4.405628422840535e-15
x_conv2 MSE: 1.702309083626008e-13
x_max_pool2 MSE: 1.1890930406120714e-13
x.view(-1, 16 * 5 * 5) MSE: 1.1890930406120714e-13
x_F.relu MSE: 1.702051179034228e-13
x_self.fc3 MSE: 6.125766669091981e-13
x_F.log_softmax MSE: 5.059064061771479e-13
和在GPU上训练的结果做个对比:
ini
torch_model.forward = MethodType(new_forward, torch_model)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dummy_input = torch.randn(1, 1, 32, 32)
dummy_input = dummy_input.to(device)
torch.onnx.export(torch_model, dummy_input, './debug_models/before_debug.onnx', input_names=['input'])
debugger.extract_debug_model('./debug_models/before_debug.onnx', './debug_models/after_debug.onnx')
dummy_input = dummy_input.to('cpu')
debugger.run_debug_model({'input': dummy_input.numpy()}, './debug_models/after_debug.onnx')
print('print debug result')
debugger.print_debug_result()
不同地方的代码如上所示,主要涉及CPU和GPU之间的迁移
yaml
print debug result
x_conv1 MSE: 0.0
x_relu MSE: 0.0
x_max_pool1 MSE: 0.0
x_conv2 MSE: 1.0748011902705912e-07
x_max_pool2 MSE: 1.1210126160676737e-07
x.view(-1, 16 * 5 * 5) MSE: 1.1210126160676737e-07
x_F.relu MSE: 7.022121906175016e-08
x_self.fc3 MSE: 1.210780880001039e-07
x_F.log_softmax MSE: 1.2211985733756592e-07
第二次跑
yaml
print debug result
x_conv1 MSE: 0.0
x_relu MSE: 0.0
x_max_pool1 MSE: 0.0
x_conv2 MSE: 1.5317392865199508e-07
x_max_pool2 MSE: 1.493277039799068e-07
x.view(-1, 16 * 5 * 5) MSE: 1.493277039799068e-07
x_F.relu MSE: 2.7951699621553416e-07
x_self.fc3 MSE: 8.633198262941733e-07
x_F.log_softmax MSE: 9.556448503644788e-07
----------------------------------------------------------
print debug result
x_conv1 MSE: 0.0
x_relu MSE: 0.0
x_max_pool1 MSE: 0.0
x_conv2 MSE: 1.332095820316681e-07
x_max_pool2 MSE: 1.3874196724827925e-07
x.view(-1, 16 * 5 * 5) MSE: 1.3874196724827925e-07
x_F.relu MSE: 1.031602891998773e-07
x_self.fc3 MSE: 1.6803554103717033e-07
x_F.log_softmax MSE: 7.133683652682521e-07
可能的优化措施:
1、自动根据 MSE 确定误差范围,
2、张量
如何解决转onnx中的精度丢失问题
通过工具定位哪一层,然后深入到层内继续断点调试。
参考文章:
zhuanlan.zhihu.com/p/516920606
zhuanlan.zhihu.com/p/489090393
zhuanlan.zhihu.com/p/513387413
zhuanlan.zhihu.com/p/493955209
zhuanlan.zhihu.com/p/498425043
建议和不足:
留给大家