hola** ONNX计算图修改神器
做过部署的小伙伴都知道，在利用TensorRT部署到NVIDIA显卡上时，onnx模型的计算图不好修改，而看了人家NCNN开发者nihui大佬的操作就知道，很多时候大佬是将onnx转换成ncnn的.paran和.bin文件后对.param的计算图做调整的，看的我心痒痒，就想有没有一种工具可以修改onnx计算图，这样，我可以合并op后，自己写个TRT插件就好了嘛
安装onnx_graphsurgeon
在新版本的TensoRT预编译包里有.whl的python包直接安装就可以了，笔者今天主要是讲怎么用这个工具，以官方的例子很好理解
生成一个onnx计算图
import onnx_graphsurgeon as gs
import numpy as np
import onnx
# Register functions to make graph generation easier
@gs.Graph.register()
def min(self, *args):
return self.layer(op="Min", inputs=args, outputs=["min_out"])[0]
@gs.Graph.register()
def max(self, *args):
return self.layer(op="Max", inputs=args, outputs=["max_out"])[0]
@gs.Graph.register()
def identity(self, inp):
return self.layer(op="Identity", inputs=[inp], outputs=["identity_out"])[0]
# Generate the graph
graph = gs.Graph()
graph.inputs = [gs.Variable("input", shape=(4, 4), dtype=np.float32)]
# Clip values to [0, 6]
MIN_VAL = np.array(0, np.float32)
MAX_VAL = np.array(6, np.float32)
# Add identity nodes to make the graph structure a bit more interesting
inp = graph.identity(graph.inputs[0])
max_out = graph.max(graph.min(inp, MAX_VAL), MIN_VAL)
graph.outputs = [graph.identity(max_out), ]
# Graph outputs must include dtype information
graph.outputs[0].to_variable(dtype=np.float32, shape=(4, 4))
onnx.save(gs.export_onnx(graph), "model.onnx")
这就是Clip操作嘛
现在就是想使用onnx_graphsurgeon这个工具将OP Min和Max整合成一个叫Clip的心OP这样即使部署时也只需要写个Clip插件就好了，当然本文只是为了演示，Clip OP已经TensorRT支持了。
修改开始
方法非常简单，先把你想要合并的OP和外界所有联系切断，然后替换成新的ONNX OP保存就好了
还不理解？上才艺
就是把Min和Identity断开，Min和c2常数断开，Max和c5常数断开，Max和下面那个Identity断开，然后替换成新的OP就好
看代码
import onnx_graphsurgeon as gs
import numpy as np
import onnx
# 这里写成函数是为了，万一还需要这样的替换操作就可以重复利用了
@gs.Graph.register()
def replace_with_clip(self, inputs, outputs):
# Disconnect output nodes of all input tensors
for inp in inputs:
inp.outputs.clear()
# Disconnet input nodes of all output tensors
for out in outputs:
out.inputs.clear()
# Insert the new node.
return self.layer(op="Clip", inputs=inputs, outputs=outputs)
# Now we'll do the actual replacement
# 导入onnx模型
graph = gs.import_onnx(onnx.load("model.onnx"))
tmap = graph.tensors()
# You can figure out the input and output tensors using Netron. In our case:
# Inputs: [inp, MIN_VAL, MAX_VAL]
# Outputs: [max_out]
# 子图的需要断开的输入name和子图需要断开的输出name
inputs = [tmap["identity_out_0"], tmap["onnx_graphsurgeon_constant_5"], tmap["onnx_graphsurgeon_constant_2"]]
outputs = [tmap["max_out_6"]]
# 断开并替换成新的名叫Clip的 OP
graph.replace_with_clip(inputs, outputs)
# 删除现在游离的子图
graph.cleanup().toposort()
# That's it!
onnx.save(gs.export_onnx(graph), "replaced.onnx")
完成onnx计算图修改
开发模版
import onnx_graphsurgeon as gs
import argparse
import onnx
import numpy as np
import json
def process_graph(graph):
node = None
for node in graph.nodes:
if node.name == "/image_encoder/patch_embed/proj/Conv":
input_tensor = gs.Variable(name="image", dtype=np.float32, shape=(3, 1024, 1024))
node.inputs[0] = input_tensor
graph.inputs = [input_tensor]
return graph
def main():
parser = argparse.ArgumentParser(description="Modify DCNv2 plugin node into ONNX model")
parser.add_argument("-i", "--input",
help="Modify ONNX Model Graph",
default="models/centertrack_DCNv2_named.onnx")
parser.add_argument("-o", "--output",
help="Path to output ONNX model with 'DCNv2_TRT' node",
default="models/modified.onnx")
args, _ = parser.parse_known_args()
graph = gs.import_onnx(onnx.load(args.input))
graph = process_graph(graph)
# 删除现在游离的子图
graph.cleanup().toposort()
onnx.save(gs.export_onnx(graph), args.output)
if __name__ == '__main__':
main()
onnx_graphsurgeon的知识
onnx_graphsurgeon只有三个ir表示，Graph，Node，Tensor
Graph
"""
Args:
nodes (Sequence[Node]): 图中nodes，一个list[Node]
inputs (Sequence[Tensor]): 图中输入tensors，一个list[Tensor]
outputs (Sequence[Tensor]): 图中输出tensors，一个list[Tensor]
name (str): 图名称，默认是"onnx_graphsurgeon_graph"
doc_string (str): 图的doc描述，默认是空字符串""
opset (int): opset版本
"""
# 获取graph
graph = gs.import_onnx(onnx.load(onnx-model.onnx))
# 主要API
#从图中删除未使用的节点和张量。
cleanup()
# 对图形进行拓扑排序。
toposort(recurse_subgraphs=True)
# 获取所有tensors
tensors(check_duplicates=False)
# 做常量折叠，在做这个操作之前必须调用toposort
fold_constants(fold_shapes=True, recurse_subgraphs=True, partitioning=None, error_ok=True)
# 创建一个节点，将其添加到此图中，并可选择创建其输入和输出张量。
layer(self, inputs=[], outputs=[], *args, **kwargs):
# 一般使用方法
@gs.Graph.register()
def add(self, a, b):
return self.layer(op="Add", inputs=[a, b], outputs=["add_out_gs"])
graph.add(a, b)
Node
class Node
"""
节点表示图中的一个操作，消耗零个或多个张量，并产生零个或更多张量。
Args:
op (str): 算子的类型
name (str): 算子的名字
attrs (Dict[str, object]): 属性，一个字典类型的
inputs (List[Tensor]): 输入的Tensor
outputs (List[Tensor]): 输出的Tensor
"""
# 获取方法
nodes = graph.nodes
Tensor
Tensor是Variable和Variable和LazyValues的基类，所以Tensor没有构造函数，一般不直接使用，所以一般使用派生类
# 判断tensor是否为空
is_empty()
"""
修改此张量以将其转换为常量。这意味着张量的所有消费者/生产者都将看到更新。
Args:
values（np.ndarray）：此张量的值
data_location（int）中的值：一个枚举值，指示存储张量数据的位置。通常，这将来自onnx.TensorProto.DataLocation。
"""
to_constant(values: np.ndarray, data_location: int = None）
"""
修改此张量以将其转换为变量。这意味着张量的所有消费者/生产者都将看到更新。
Args:
dtype（np.dtype）：张量的数据类型。
shape（Sequence[int]）：张量的形状。
"""
to_variable(self, dtype: np.dtype = None, shape: Sequence[Union[int, str]] = [])
# 获取方法
tensors = graph.tensors() # 返回Tensor的key:value字典
# 断开连接
tensors = graph.tensors()
tensor = tensors['op_name']
tensor.outputs.clear()
Variable
变量是Tensor的派生类
class Variable
"""
表示一个张量，其值在推断时间之前是未知的。
Args:
name（str）：张量的名称。
dtype（numpy.dtype）：张量的数据类型。
shape（Sequence[Unint[int，str]]）：张量的形状。如果模型使用标注参数，则可能包含字符串。
"""
# 创建方法示例
input_tensor = gs.Variable(name="image", dtype=np.float32, shape=(3, 1024, 1024))
"""
变量转成常量
"""
to_constant(values: np.ndarray)
Constant
常量是Tensor的派生类
class Constant
"""
表示值已知的张量。
Args:
name（str）：张量的名称。
values（numpy.ndarray）：这个张量中的值，以numpy数组的形式。
data_location（int）：一个枚举值，指示存储张量数据的位置。通常，这将来自onnx.TensorProto.DataLocation。
"""
to_variable(dtype: np.dtype = None, shape: Sequence[Union[int, str]] = [])
LazyValues
一个特殊的对象，它表示应该延迟加载的常量张量值。
load()
从基本张量值加载numpy数组。
