需求真的是千奇百怪，最近项目需要修改多年前通过tensorflow转换得到的ONNX模型，关键转换前的tensorflow模型已经神秘地失踪了
本小姐真是无力吐槽，这个班真是一天都不想上了，冷静下来想想，这个“不想上班”的想法还是太年轻，毕竟挣钱要紧，然后记录一波打工人的艰难之旅，dddd（懂得都懂）
友情提示：阅读该内容大概需要20分钟，请理性安排，适当的时候可以知难而退
目录
make_tensor_value_info(name,elem_type,shape)
make_node(op_type, inputs, outputs, name=None)
make_graph(nodes,name,inputs,outputs,initializer=None)
make_tensor(name,data_type,dims,vals)
一、ONNX简介
官方介绍，开放神经网络交换（Open Neural Network Exchange）简称ONNX，是微软和Facebook提出用来表示深度学习模型的开放格式。所谓开放就是ONNX定义了一组和环境，平台均无关的标准格式，来增强各种AI模型的可交互性。
换句话说，无论你使用何种训练框架训练模型（比如TensorFlow/Pytorch/…），在训练完毕后都可以将这些框架的模型统一转换为ONNX这种统一的格式进行存储。ONNX文件不仅仅存储了神经网络模型的权重，同时也存储了模型的结构信息以及网络中每一层的输入输出和一些其它的辅助信息。
二、ONNX结构分析
ONNX是把一个网络的每一层或者说一个算子当成节点 node，使用这些 Node去构建一个 Graph，即一个网络。最后将 Graph和其它的生产者信息，版本信息等合并在一起生成一个 Model，也即是最终的ONNX模型文件。在构建ONNX模型的时候，https://github.com/onnx/onnx/blob/master/onnx/helper.py
这个文件需要划重点，其中make_node、make_graph、make_model是不可或缺的。make_tensor_value_info和make_tensor是构建graph中所需要用到的
make_tensor_value_info(name,elem_type,shape)
name: 节点名字 [类型:字符串]
elem_type: 数据类型 [类型:TensorProto.DataType]
shape: 数据维度(形状) [类型:int列表/元组]
model = onnx.load("./bert.onnx")
# 给模型插入一个冗余的输入节点
ori_segment_ids = onnx.helper.make_tensor_value_info("ori_segment_ids_ph:0", 6, shape=[None, None])
model.graph.input.append(ori_segment_ids)make_node(op_type, inputs, outputs, name=None)
op_type: 节点的算子类型 [类型:字符串]，详细可以参考onnx给出的算子列表
inputs: 存放节点输入的名字 [类型:字符串列表]，每个节点输入的数量根据情况会有不同
outputs: 存放节点输出的名字 [类型:字符串列表]，与inputs类似，同样需要根据官网给出的输出个数来设置，大多数情况是一个输出
name: 节点名，可有可无，不要和op_type搞混了
# 用于扩充维度 1*100 → 1*1*100
unsqueeze_node = onnx.helper.make_node(
op_type="Unsqueeze",
inputs=["insert_unsqueeze1"],
outputs=["insert_unsqueeze"],
axes=[0],
)ONNX 对节点的输入要求：一个节点的输入，要么是整个模型的输入，要么是之前某个节点的输出。
make_graph(nodes,name,inputs,outputs,initializer=None)
nodes: 用make_node生成的节点列表 [类型:NodeProto列表]
name: graph的名字 [类型:字符串]
inputs: 存放graph的输入数据信息 [类型:ValueInfoProto列表]
输入数据的信息以ValueInfoProto的形式存储，会用到make_tensor_value_info，来将输入数据的名字、数据类型、形状(维度)给记录下来。对于一个网络而言如何能体现其网络结构呢？即节点与节点之间的关联。在构建每一个node时就需要注意，当前node的输入来自于哪一个node的输出，名字要匹配上，才能将node间联系体现出来。
outputs: 存放graph的输出数据信息 [类型:ValueInfoProto列表]，与inputs相同。
initializer: 存放超参数 [类型:TensorProto列表]，对于一个多层网络而言，其中间层的输入有来自上一层的输出，也有来自外界的超参数和数据。 initializer作为存放超参数具体数值的TensorProto列表，其中每个TensorProto总会有与其对应的ValueInfoProto存在，对应关系通过name来联系。
input1 = np.random.rand(1, 3, 4, 5).astype("float32")
input2 = np.random.rand(1, 5).astype("float32")
inputs = [helper.make_tensor_value_info("input1", TensorProto.FLOAT, shape=(1, 3, 4, 5)),
helper.make_tensor_value_info("input2", TensorProto.FLOAT, shape=(1, 5))]
outputs = [helper.make_tensor_value_info("output", TensorProto.FLOAT, shape=(1, 3, 4, 5))]
nodes = [helper.make_node("Add", ["input1", "input2"], ["output"])]
graph = helper.make_graph(nodes, "bcast_test", inputs, outputs)
bcast_model = helper.make_model(graph)其中，make_graph 的节点参数要求：计算图的节点必须以拓扑序给出。
make_model(graph, **kwargs)
graph: 用make_graph生成的GraphProto
make_tensor(name,data_type,dims,vals)
name: 数据名字，要与该数据的信息tensor value info中名字对应 [类型:字符串]
data_type: 数据类型 [类型:TensorProto.DataType] 如TensorProto.FLOAT、TensorProto.UINT8、TensorProto.FLOAT16等
dims: 数据维度 [类型:int列表/元组]
vals: 数据值，好像要可迭代的 [类型:任意]raw:选择是否用二进制编码 [类型:bool]
tensor = make_tensor('test_tensor', onnx.TensorProto.FLOAT, [1], [1])make_attribute(key,value)
key: 键值 [类型:字符串]
value: 数值 [类型:任意]
attr = helper.make_attribute("int", 5)三、ONNX模型修改
大家都知道，模型结构优化一直以来都是比较fancy的工作，太fancy我么也做不了，所以我们今天只简单地修改一下模型结构
合并两个ONNX模型
import onnx
model_1 = onnx.load(file1)
onnx.compose.expand_out_dim(model_1, dim_idx=0, inplace=True)#用于扩展维度 [None,None] [1,None,None]
model_2 = onnx.load(file2)
combined_model=onnx.compose.merge_models(model_1, model_2, io_map=[('output:0', 'input:0')]) # 分别为两个模型中待合并的节点
查看ONNX节点信息
import onnx
model = onnx.load("./bert.onnx")
outputs = model.graph.output # 输出节点
# input_node = model.graph.input # 输入节点
# nodes = model.graph.node # 所有节点
print(outputs )
＿＿＿＿＿＿＿＿＿输出信息如下＿＿＿＿＿＿＿＿＿
[name: "label_score:0"
type {
tensor_type {
elem_type: 1
shape {
dim {
dim_param: "unk__403"
}
dim {
dim_value: 2
}
}
}
}]
＿＿＿＿＿＿＿＿＿＿输出信息＿＿＿＿＿＿＿＿＿＿其中，elem_type为输出节点的类型，数据类型（elem_type）共有16种，
elem_type: 1 --> float32
elem_type: 2 --> uint8
elem_type: 3 --> int8
elem_type: 4 --> uint16
elem_type: 5 --> int16
elem_type: 6 --> int32
elem_type: 7 --> int64
elem_type: 8 --> string
elem_type: 9 --> boolean
elem_type: 10 --> float16
elem_type: 11 --> float64
elem_type: 12 --> uint32
elem_type: 14 --> uint64
elem_type: 15 --> complex128
elem_type: 16 --> bfloat16 其余的数字全为undefined
修改输入节点type和名称
model = onnx.load("./bert.onnx")
input_node = model.graph.input
for i in range(len(input_node)):
if input_node[i].name == "ori_input_quests:0":
model.graph.input[i].type.tensor_type.elem_type = 6
model.graph.input[0].name = "proj/cond_1/Merge:0" #修改输入节点名字
nodes = model.graph.node
length = len(nodes)
for i in range(length):
if nodes[i].name == "cond_If__38":
model.graph.node[i].output[0] = "proj/cond_1/Merge:0" # 修改节点名字插入Cast节点/算子
通过算子原型构建Graph时，要求前后算子的dtype必须一致，上一个算子的输出dtype如果和下一层算子的输入dtype不匹配时，需要插入Cast算子
model = onnx.load("./bert.onnx")
nodes = model.graph.node
for i in range(len(nodes)):
if nodes[i].output[0] == "_not_equal_scalar0_eq":
index = i # 记录节点id
break
else:
raise ValueError("find insert_node error!!!")
cast_node = onnx.helper.make_node(
op_type="Cast",
inputs=["ori_input_quests:0"],
outputs=["insert_cast_0"],
to=getattr(onnx.TensorProto, "FLOAT")
)
model.graph.node.insert(index, cast_node) # 插入节点
model.graph.node[index + 1].input[0] = "insert_cast_0" # 指向插入的节点size算子
将张量作为输入并输出一个 int64 标量，该标量等于输入张量的元素总数
size_node = onnx.helper.make_node(
op_type="Size",
inputs=["ori_input_quests:0"],
outputs=["insert_Size"],
)插入range算子
Range算子类似python的range()函数
model = onnx.load("./bert.onnx")
node = model.graph.node
for i in range(len(node)):
if node[i].output[0] == "bert/embeddings/position_embeddings_indices_casted":
index = i # 搜索节点所在的网络id
break
else:
raise ValueError("查找节点失败")
initializer = model.graph.initializer # ONNX初始化值
start_const = onnx.helper.make_tensor(name='start_const',
data_type=6,
dims=[],
vals=[0])
# 初始化 常量 0和1
step_const = onnx.helper.make_tensor(name='step_const',
data_type=6,
dims=[],
vals=[1])
# 将生成的tensor插入途中
model.graph.initializer.append(start_const)
model.graph.initializer.append(step_const)
range_node = onnx.helper.make_node(
op_type="Range",
inputs=["start_const", "insert_Size", "step_const"],
outputs=["src_position"],
)
# 若"insert_Size"=10，则range之后生成[0,1,2,3,4,5,6,7,8,9]
model.graph.node.insert(index, range_node)
model.graph.node[index + 1].input[0] = "src_position" # 插入节点指向下一个节点增加、删除输入输出节点
model = onnx.load("./bert.onnx")
# 插入一个输入节点
ori_segment_ids = onnx.helper.make_tensor_value_info("ori_segment_ids_ph:0", 6, shape=[None, None])
model.graph.input.append(ori_segment_ids)
# 删除
input_node = model.graph.input
for i in range(len(input_node)):
if input_node[i].name == "src_positions":
model.graph.input.remove(input_node[i])
break
onnx.checker.check_model(model) # 检测模型格式
onnx.save(model, "./new.onnx") # 保存模型pad算子
有时候，已有的模型输入维度固定的，比如[1,100]，但我们 想让输入支持可变长度的[1,None]，需要用到padding OP
实现方法：支持输入动态，而且模型里面的操作都是固定维度，需要先修改输入维度为动态，然后进步模型后在修改维度到固定值；
model = onnx.load("./bert.onnx")
inputs = model.graph.input
for i in range(len(inputs)):
if inputs[i].name == "ori_input_quests:0":
model.graph.input[i].type.tensor_type.elem_type = 6 # 修改输入类型
model.graph.input[i].type.tensor_type.shape.dim[1].dim_param = "len" # 修改输入维度为动态长度
input_len = 100 # 假设原始的输入长度 [1,100]
pads_value = onnx.helper.make_tensor(name='pads_value',
data_type=6,
dims=[],
vals=[0])
pads_const = onnx.helper.make_tensor(name='pads_const',
data_type=7,
dims=[4],
vals=[0, 0, 0, input_len])
# 用于padding, 结束位置补100个0 [1,200] vals：第一个维度左边补0个数，第二个维度左边补0个数，第一个维度右边补0个数，第二个维度右边补0个数
# data =[
# [1.0, 1.2],
# [2.3, 3.4],
# [4.5, 5.7],]
# pads = [0, 2, 0, 0]
# mode = 'constant'
# constant_value = 0.0
# output =[
# [0.0, 0.0, 1.0, 1.2],
# [0.0, 0.0, 2.3, 3.4],
# [0.0, 0.0, 4.5, 5.7],]
start_const = onnx.helper.make_tensor(name='start_const',
data_type=7,
dims=[2],
vals=[0, 0])
end_const = onnx.helper.make_tensor(name='end_const',
data_type=7,
dims=[2],
vals=[1, input_len])
# 用于切片操作, 第一个维度[0:1] 第二个维度[0:100]
model.graph.initializer.append(pads_value)
model.graph.initializer.append(pads_const)
model.graph.initializer.append(start_const)
model.graph.initializer.append(end_const)
pad_node = onnx.helper.make_node(
op_type="Pad",
inputs=["ori_input_quests:0", "pads_const", "pads_value"],
outputs=["insert_pad"],
mode='constant',
)
slice_node = onnx.helper.make_node(
"Slice",
inputs=["insert_pad", "start_const", "end_const"],
outputs=["insert_slice"],
)
nodes = model.graph.node
for i in range(len(nodes)):
if nodes[i].output[0] == "bert/embeddings/word_embeddings_indices_casted":
index = i
break
else:
raise ValueError("查找节点失败")
model.graph.node.insert(index, pad_node)
model.graph.node.insert(index + 1, slice_node)
model.graph.node[index + 2].input[0] = "insert_slice"
onnx.checker.check_model(model) # 检测模型格式
from onnx import shape_inference
model = shape_inference.infer_shapes(model) # 维度推断
onnx.save(model, "./new.onnx") # 保存模型插入Slice算子
类似python的切片操作
model = onnx.load("./bert.onnx")
start_const = onnx.helper.make_tensor(name='start_const',
data_type=7,
dims=[2],
vals=[0, 0])
end_const_1 = onnx.helper.make_tensor(name='end_const_1',
data_type=7,
dims=[2],
vals=[1, -1])
# 用于切片操作，第一个维度[0:1] 第二个维度[0:-1] eg:维度 1*15 → 1*14
model.graph.initializer.append(start_const)
model.graph.initializer.append(end_const_1)
slice_node = onnx.helper.make_node(
"Slice",
inputs=["ori_input_quests:0", "start_const", "end_const_1"],
outputs=["insert_slice_2"],
)
nodes = model.graph.node
for i in range(len(nodes)):
if nodes[i].output[0] == "insert_pad":
model.graph.node.insert(i, slice_node)
model.graph.node[i + 1].input[0] = "insert_slice_2"
breakSqueeze算子
维度1*1*100 需要 变为 1*100，可以用到Squeeze，相反的操作就是Unsqueeze，用于扩充维度
squeeze_node_1 = onnx.helper.make_node(
op_type="Squeeze",
inputs=["ori_input_quests:0"],
outputs=["insert_Squeeze_1"],
axes=[0],
)shape算子
获取维度
size_node_1 = onnx.helper.make_node(
op_type="Shape",
inputs=["insert_Squeeze_1"],
outputs=["insert_size_1"],
)
# torch.tensor([[2, 3, 4], [3,4,5]])
# shape操作输出 (2, 3)Sub算子
model = onnx.load("./bert.onnx")
sub_const = onnx.helper.make_tensor(name='sub_const',
data_type=7,
dims=[1],
vals=[2])
model.graph.initializer.append(sub_const)
sub_contant = onnx.helper.make_node(
"Sub",
inputs=["insert_size_1", "sub_const"],
outputs=["insert_sub"],
)
Concat算子
将list of tensor合并成一个tensor
contact_node = onnx.helper.make_node(
op_type="Concat",
inputs=["insert_size_2", "insert_sub"],
outputs=["insert_contact"],
axis=0,
)ConstantOfShape算子
生成具有给定值和形状的张量
constant_node = onnx.helper.make_node(
"ConstantOfShape",
inputs=["insert_contact"],
outputs=["insert_contant"],
value=onnx.helper.make_tensor("value", onnx.TensorProto.FLOAT, [1], [1]),
)
# 根据inputs的维度生成全1的向量，eg[[1,1],[1,1]注意：
如果新插入了一个节点，没有删除相应的节点，想要在保存模型checkmodel的时候通过就必须要保证，插入节点的insert位置就是在原先节点的位置，然后在输出节点之后的名字再进行修改，输入节点也要进行修改，就是onnx模型中的输入节点一定要在该节点之前存在，否则就会导致拓扑不可分，就会报错
【参考】
ONNX 算子文档：https://github.com/onnx/onnx/blob/main/docs/Operators.md#Unsqueeze
onnx源码：https://github.com/onnx/onnx/tree/main/onnx
ONNX常用函数：https://github.com/onnx/onnx/blob/main/docs/PythonAPIOverview.md
2384
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
