前言
杜老师推出的 tensorRT从零起步高性能部署 课程，之前有看过一遍，但是没有做笔记，很多东西也忘了。这次重新撸一遍，顺便记记笔记。
本次课程学习 tensorRT 基础-实际模型上 onnx 文件的各种操作
课程大纲可看下面的思维导图
1. onnx
本节课程我们来学习 onnx 的补充，拿到一个实际的 onnx 后如何读取到你想要的东西
本次演示使用的是 yolov5s.onnx，这并不重要，可以任意拿一个 onnx，我们的目的是学习如何读取实际的 onnx
1.1 读取节点
我们先来获取第一个卷积 model.0.conv.weight 的权重信息
import onnx
import numpy as np
model = onnx.load("yolov5s.onnx")
for item in model.graph.initializer:
if item.name == "model.0.conv.weight":
print("shape: ", item.dims)
weight = np.frombuffer(item.raw_data, dtype=np.float32).reshape(32, 3, 6, 6)
print(weight)
运行效果如下：
上节课有提到 conv 的权重存储在 initiallizer 中，因此我们可以遍历然后通过 name 进行筛选，获取我们想要得到的 model.0.conv.weight，通过 numpy 将 raw_data 数据转换成 numpy 格式的 float32 的数据，从获取的权重和 onnx 模型的权重对比可看出二者都一一对应，说明我们获取的信息没问题。
1.2 修改节点
我们再来看下 Add 节点中的 Constant 类型：
import onnx
import numpy as np
for item in model.graph.node:
if item.op_type == "Constant":
if '357' in item.output:
t = item.attribute[0].t
data = np.frombuffer(t.raw_data, np.float32).reshape(*t.dims)
print(data.shape)
print(data)
前面我们提到过对于 anchor grid 类的常量数据，通常会储存在 model.graph.node 中，并指定类型为 Constant，因此我们遍历了 node 节点，通过 name 来筛选，筛选之后打印的 item 数据量非常大，看不出我们想要的东西，我们可以通过调试进一步观察：
可以看到 item 有一个 attribute 方法，它是一个数组，且长度为 1，我们通过访问 item.attribute 可以看到对应的 name、dims、raw_data 信息了，我们可以进一步调试：
可以看到最终我们想要获取的数据可以通过 item.attribute[0].t.raw_data 获得
可以发现我们最终的 shape 和数据信息都可以对应上，这是关于 constant 的读取
接下来我们尝试修改下 Constant 的值，修改的是 Slice 的 start：
import onnx
import numpy as np
model = onnx.load("yolov5s.onnx")
for item in model.graph.node:
if item.op_type == "Constant":
if '373' in item.output:
t = item.attribute[0].t
print(type(t))
print(t)
print(np.frombuffer(t.raw_data, dtype=np.int64))
t.raw_data = np.array([666], dtype=np.int64).tobytes()
onnx.save(model, "new.onnx")
我们先获取了对应的 raw_data 数据，然后将其修改为我们想要的值，运行效果如下：
从图中可以看出修改后导出的 onnx 模型的 Constant 的值已经修改为了 666
Slice 中的 Constant 数据类型并不是 float32 而是 int64，我们打印其对应的 data_type 显示数字为 7，然后去 onnx-ml.proto 中可以找到 message TensorProto 可以查询数字 7 对应的数据类型是什么
1.3 替换节点
这次我们来个高级点的，替换 reshape 节点：
import onnx
import numpy as np
import onnx.helper as helper
model = onnx.load("yolov5s.onnx")
for item in model.graph.node:
if item.name == "Reshape_242":
# print(item)
new_reshape = helper.make_node("Reshape", ["377", "378"], "379", "Reshape_666", myname="Zhou")
item.CopyFrom(new_reshape)
onnx.save(model, "new.onnx")
导出的 onnx 前后对比图如下：
可以看到修改后导出的 onnx 确实是按照我们想要的都修改了，同时还添加了新的属性。值得注意的是，我们在替换节点时不能直接 = 替换，如果想要替换节点，你需要调用 protobuf 提供的函数 CopyFrom
1.4 删除节点
我们接下来演示下如何删除某个节点，以 Transpose 节点为例：
import onnx
import numpy as np
import onnx.helper as helper
model = onnx.load("yolov5s.onnx")
find_node_with_input = lambda name: [item for item in model.graph.node if name in item.input][0]
find_node_with_output = lambda name: [item for item in model.graph.node if name in item.output][0]
remove_nodes = []
for item in model.graph.node:
if item.name == "Transpose_201":
# 上一个节点的输出是当前节点的输入
prev = find_node_with_output(item.input[0])
# 下一个节点的输入是当前节点的输出
next = find_node_with_input(item.output[0])
next.input[0] = prev.output[0]
remove_nodes.append(item)
for item in remove_nodes[::-1]:
model.graph.node.remove(item)
onnx.save(model, "new.onnx")
我们在删除了某个节点后，你需要将 input 和 output 对接起来，不能让它们直接断开。因此需要将删除节点的上个节点输出和删除节点的下个节点的输入拼接，然后再删除。删除的时候也要注意，最好不要边循环边删除，应该分为两部分来做这个事情
1.5 修改input和output
我们同样可以将输入和输出的动态 batch 修改为静态 batch，代码如下：
import onnx
import numpy as np
import onnx.helper as helper
model = onnx.load("yolov5s.onnx")
static_batch_size = 3
input = model.graph.input[0]
print(type(input))
new_input = helper.make_tensor_value_info("images", 1, [static_batch_size, 3, 640, 640])
print(new_input)
model.graph.input[0].CopyFrom(new_input)
new_output = helper.make_tensor_value_info("output", 1, [static_batch_size, 25200, 7])
model.graph.output[0].CopyFrom(new_output)
onnx.save(model, "new.onnx")
运行效果如下：
我们可以打印 input 节点的信息方便后续我们修改为静态 batch，可以看到其 name 为 images，elem_type 为 1，代表 float32。下图对比了修改前后的 onnx，可以看到确实 input 和 output 发生了修改
同样的，你想静态的修改为动态的，做法也是一样的，因此虽然你的 onnx 导出来是一个确定的东西，但是后期如果你想要修改它，随便改！！！
1.6 预处理的接入
我们同样可以将一些复杂的预处理塞到 onnx 中，其代码如下：
import onnx
import numpy as np
import onnx.helper as helper
model = onnx.load("yolov5s.onnx")
import torch
class Preprocess(torch.nn.Module):
def __init__(self) -> None:
super().__init__()
self.mean = torch.rand(1, 1, 1, 3)
self.std = torch.rand(1, 1, 1, 3)
def forward(self, x):
# x = B x H x W x C Uint8
# y = B x C x H x W Float32 减均值除标准差
x = x.float()
x = (x / 255.0 - self.mean) / self.std
x = x.permute(0, 3, 1, 2)
return x
pre = Preprocess()
torch.onnx.export(
pre, (torch.zeros(1, 640, 640, 3, dtype=torch.uint8)), "preprocess.onnx"
)
pre_onnx = onnx.load("preprocess.onnx")
# 1. 先把 pre_onnx 的所有节点以及输入输出名称加上前缀，防止命名冲突
# 2. 把 yolov5s.image 为输入的节点，修改为 pre_onnx 的输出节点
# 3. 把 pre_onnx 的 node 全部放到 yolov5s 的 node 中
# 4. 把 pre_onnx 的输入作为 yolov5s 的 input名称
for n in pre_onnx.graph.node:
n.name = f"pre_{
n.name}"
for i in range(len(n.input)):
n.input[i] = f"pre_{
n.input[i]}"
for i in range(len(n.output)):
n.output[i] = f"pre_{
n.output[i]}"
for n in model.graph.node:
if n.name == "Conv_0":
n.input[0] = "pre_" + pre_onnx.graph.output[0].name
for n in pre_onnx.graph.node:
model.graph.node.append(n)
input_name = "pre_" + pre_onnx.graph.input[0].name
model.graph.input[0].CopyFrom(pre_onnx.graph.input[0])
model.graph.input[0].name = input_name
onnx.save(model, "new.onnx")
首先，我们导出了一个预处理的 onnx，如下图所示：
然后我们需要考虑如何将两个 onnx 对接起来，主要步骤如下：
1、 先把 pre_onnx 的所有节点以及输入输出名称加上前缀，防止命名冲突
2、 把 yolov5s.image 为输入的节点，修改为 pre_onnx 的输出节点
3、 把 pre_onnx 的 node 全部放到 yolov5s 的 node 中
4、 把 pre_onnx 的输入作为 yolov5s 的 input名称
拼接后的 onnx 如下图所示：
当你的算子很难搞定时，可以考虑 pytorch 生成下来然后塞进去，这样灵活度就很高了。预处理可以接入，同样后处理也可以这样接入，所以你可以把 onnx 当成一个类似可编辑的记事本一样的东西就行，因此无论你遇到多复杂的 onnx，你都可以去编辑修改它，而不至于束手无策。
总结
本节课程我们以实际的 yolov5s.onnx 模型为例，学习了 onnx 文件的各种操作，包括基本的读取某个节点的信息，修改某个节点的值，高级一点的包括替换对应的节点，删除对应的节点，还有将动态 shape 修改为 静态 shape，最后我们还尝试在原有s onnx 的基础上接了一段预处理的代码。
通过本次课程的学习，我们需要知道对于 onnx 而言，我们是可以操作的，读取、修改、删除、增加等等，你把它当成一个可编辑的东西就行，后续无论是多么复杂的 onnx，你都应该要知道如何处理。