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數組。