onnx模型图优化/模型修改

比较怪异的是个别模型部分initializer也是graph.input的一部分，这是因为把部分initializer的信息也添加到了graph的input里面，这可能导致onnx-simplifier简化模型出错。解决方法是把它们从input里面删掉即可：找到这些invalid input name的idx，从大到小排好序后一次del graph.input[idx]即可。

保存模型：

onnx.save(onnx_model, model_path_out, save_as_external_data=False)
save_as_external_data设置为True用于处理大于2GB的模型，这会把权重单独保存到一个文件。

手动序列化：

with open("model.onnx", "wb") as f:
f.write(onnx_model.SerializeToString())

model的各个组件可以单独调用SerializeToString，比如mode.graph，或者graph内部的node。

打印出onnx_model.graph.node可以使用的方法(类似地可以打印其他类包含的方法)：

for item in dir(onnx_model.graph.node):
print(item)
可以看到跟python list基本一样，可用的方法为：add, append, extend, insert, pop, remove, reverse, sort

大多数proto都是以list的方式存放的，比较遗憾没有一个clear的方法来清除所有内容，可以通过循环的方式来删除，例如：

for i in reversed(range(len(graph.value_info))):
    del graph.value_info[i]

获取onnx模型opset:

print("domain:", onnx_model.opset_import[0].domain)
print("opset version:", onnx_model.opset_import[0].version)

修改算子属性attrs

修改node attr一个最简单方法是删除旧的，添加个新的。

node添加attr： node.attribute.append(onnx.helper.make_attribute("axes", [-1]))

可以给onnx标准的node添加用户自定义的attr并且可以正常onnxsim和模型检查，但注意自定义的attr名称需要双下划线开头，如__custom_attr。

创建简单模型

tensorflow 的graph_def的node输出名称不用指定，默认是node_name:idx形式，没有idx则默认为第0个输出。

而onnx模型的node的output name可以不是node_name:idx的形式，可以是任意的有效字符串。

这是因为onnx的tensor和node是分离的，每个tensor可以有一个任意独特的名称，而然后把tensor的名称赋予node的input和output, graph的initializer，input output作为输入输出。

用onnx原生创建一个包含layernorm算子的图（使用onnx_graphsurgeon工具创建模型图更简单）：

import onnx
import numpy as np
from onnx.helper import make_node, make_graph, make_tensor, make_model, make_opsetid

inputs = [onnx.helper.make_tensor_value_info(name="input", elem_type=onnx.TensorProto.FLOAT, shape=(32, 512))]

outputs = [onnx.helper.make_tensor_value_info(name="output", elem_type=onnx.TensorProto.FLOAT, shape=(32, 512))]

bias_shape = [512]
gain_shape = [512]
bias_values = np.random.uniform(-10, 10, size=bias_shape).astype("float32")
gain_values = np.random.uniform(-10, 10, size=gain_shape).astype("float32")

bias_const = make_tensor(name="W", data_type=onnx.TensorProto.FLOAT, dims=bias_shape, vals=bias_values, raw=False)
gain_const = make_tensor(name="B", data_type=onnx.TensorProto.FLOAT, dims=gain_shape, vals=gain_values, raw=False)

ln_node = onnx.helper.make_node(
    "LayerNormalization",
    inputs=["input", "W", "B"],
    outputs=["output"],
    axis=-1,
    epsilon=1e-05)

# Nodes in a graph must be topologically sorted
nodes = [
    ln_node
]
initializer = [bias_const, gain_const]
graph = onnx.helper.make_graph(nodes=nodes, name="layer_norm_graph", inputs=inputs,
                               outputs=outputs, initializer=initializer)

onnx_model = onnx.helper.make_model(graph, opset_imports=[make_opsetid(domain="", version=17)])

onnx.checker.check_model(onnx_model)
onnx.save(onnx_model, "layernorm1.onnx")

创建一个包含两个add算子的图：

import onnx
import numpy as np
from onnx.helper import make_node, make_graph, make_tensor_value_info, make_model, make_opsetid

# use -1 for dynamic shape
inputs = [onnx.helper.make_tensor_value_info(name="input1", elem_type=onnx.TensorProto.FLOAT, shape=(-1, 16)),
          onnx.helper.make_tensor_value_info(name="input2", elem_type=onnx.TensorProto.FLOAT, shape=(-1, 16))]

outputs = [onnx.helper.make_tensor_value_info(name="output", elem_type=onnx.TensorProto.FLOAT, shape=(-1, 16))]

const_shape = (1, 16)
const_values = np.random.uniform(-10, 10, size=const_shape).astype("float32")

const_node0 = onnx.helper.make_node(
    op_type="Constant",
    inputs=[],
    outputs=["const1:0"],
    name="const1",
    value=onnx.helper.make_tensor(name='const1',
                                  data_type=onnx.TensorProto.FLOAT,
                                  dims=const_values.shape,
                                  vals=const_values.reshape(-1)))

add_node0 = onnx.helper.make_node(op_type="Add", inputs=["input1", "input2"], outputs=["add1:0"], name="add1")
add_node1 = onnx.helper.make_node(op_type="Add", inputs=["add1:0", "const1:0"], outputs=["output"], name="add2")

# Nodes in a graph must be topologically sorted
nodes = [
    const_node0,
    add_node0,
    add_node1,
]

graph = onnx.helper.make_graph(nodes=nodes, name="add_test", inputs=inputs, outputs=outputs)
# you can also use graph.node.insert(idx, add_node0) to insert add_node0 before node at idx

onnx_model = onnx.helper.make_model(graph, opset_imports=[make_opsetid(domain="", version=11)])

onnx.checker.check_model(onnx_model)
onnx.save(onnx_model, "add_model.onnx")

创建包含自定义domain的模型

import onnx
import numpy as np
from onnx.helper import make_node, make_graph, make_tensor_value_info, make_model, make_opsetid

inputs = [
    onnx.helper.make_tensor_value_info(name="X", elem_type=onnx.TensorProto.FLOAT, shape=(1, 1, 4, 4)),
    onnx.helper.make_tensor_value_info(name="Grid", elem_type=onnx.TensorProto.FLOAT, shape=(1, 6, 6, 2)),
]

outputs = [onnx.helper.make_tensor_value_info(name="output", elem_type=onnx.TensorProto.FLOAT, shape=(1, 1, 6, 6))]

node = onnx.helper.make_node(
    "GridSample",
    inputs=["X", "Grid"],
    outputs=["Y"],
    mode="bilinear",
    padding_mode="zeros",
    align_corners=0,
    domain="com.microsoft",
)

add_node0 = onnx.helper.make_node(op_type="Add", inputs=["Y", "Y"], outputs=["output"], name="add1")

nodes = [
    node,
    add_node0,
]

graph = onnx.helper.make_graph(nodes=nodes, name="add_test", inputs=inputs, outputs=outputs)

onnx_model = onnx.helper.make_model(graph, opset_imports=[make_opsetid(
    domain="", version=11), make_opsetid(domain="com.microsoft", version=1)])

onnx.save(onnx_model, "grid_sample.onnx")

onnx.checker.check_model(onnx_model)

onnx_graphsurgeon创建模型样例参考：

https://github.com/leimao/TensorRT-Custom-Plugin-Example/blob/main/scripts/create_identity_neural_network.py

pytorch导出onnx自定义算子

方式1：

https://github.com/leimao/TensorRT-Custom-Plugin-Example/blob/main/scripts/export_identity_neural_network_new_opset.py

非常简单，只需要使用export_modules_as_functions包含作为自定义算子的nn.Module，算子的attr, weight都能被代入，opset需要>=15。另外自定义算子的nn.Module的attr定义参考上面的案例。这种方式导出的onnx的自定义算子其对应的子图放到了onnx_model的functions里面，可以考虑进行删除。

    torch.onnx.export(model=identity_neural_network,
                      args=(input_data, ),
                      f=onnx_file_path,
                      input_names=["X0"],
                      output_names=["X3"],
                      opset_version=opset_version,
                      export_modules_as_functions={IdentityConv})

方式2 :基于symbolic方法

https://github.com/open-mmlab/mmdeploy/blob/main/docs/zh_cn/tutorial/04_onnx_custom_op.md

使用Pytorch导出自定义ONNX算子_自定义onnx模块-CSDN博客

https://github.com/Shank2358/DCNv2/blob/main/dcn_v2_onnx.py

onnx模型仓库和下载

1. transformers

在transformers网站里面直接搜索onnx模型，或者权重转模型

例如直接转好的onnx模型：

https://huggingface.co/rocca/openai-clip-js/tree/main

https://huggingface.co/philschmid/distilbert-onnx/tree/main

可以直接命令行下载和转换模型，例如：

python -m transformers.onnx --model=microsoft/swin-tiny-patch4-window7-224 onnx/ --opset 15

onnx/是本地文件路面，第一个是transformers仓库里面的模型路径

不成功时可以尝试手动下载权重和配置文件转模型，例如：

https://huggingface.co/apple/mobilevit-small/tree/main

下载上面链接的权重和配置文件保持在一个文件夹，如model

再通过下面命令转成onnx模型(需要先安装pytorch等依赖包)：

python -m transformers.onnx --model=in_model_dir out_model_dir

参考：https://huggingface.co/docs/transformers/serialization

通过git下载huggingface模型：

apt-get install git-lfs
git lfs install
git clone https://huggingface.co/${username}/${model_name}

例如下载runwayml/stable-diffusion-v1-5的onnx分支：

git clone -b onnx https://huggingface.co/runwayml/stable-diffusion-v1-5

onnx tensor proto与Numpy array转换

from onnx import numpy_helper 
np_tensor = numpy_helper.to_array(init)

numpy tensor到onnx tensor proto: make_tensor

shape infer

onnx shape推导

onnx/ShapeInference.md at main · onnx/onnx · GitHub

onnx/PythonAPIOverview.md at main · onnx/onnx · GitHub

from onnx import helper, shape_inference

        if onnx_graph.ByteSize() > 2147483648:
            temp_dir = tempfile.TemporaryDirectory().name
            os.makedirs(temp_dir, exist_ok=True)
            onnx_orig_path = os.path.join(temp_dir, 'model.onnx')
            onnx_inferred_path = os.path.join(temp_dir, 'inferred.onnx')
            onnx.save_model(onnx_graph,
                onnx_orig_path,
                save_as_external_data=True,
                all_tensors_to_one_file=True,
                convert_attribute=False)
            onnx.shape_inference.infer_shapes_path(onnx_orig_path, onnx_inferred_path)
            onnx_graph = onnx.load(onnx_inferred_path)
        else:
            onnx_graph = shape_inference.infer_shapes(onnx_graph)

也可以借助用onnx-simplifier来进行shape infer。在原有的shape不全或者部分错误的时候可能报错。

onnxruntime里面也有个shape infer工具：from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference。但是不太建议，有些bug。

包含自定义domain的op shape infer:

先用上面的工具infer shape，然后自己给出自定义算子的shape，然后再调用上面的工具infer shape。

得到的shape存在graph的value_info

value_info {
  name: "MatMul_89"
  type {
    tensor_type {
      elem_type: 1
      shape {
        dim {
          dim_value: 1
        }
        dim {
          dim_value: 40
        }
        dim {
          dim_value: 40
        }
      }
    }
  }
}

获取onnx每个tensor的shape (存放在graph.value_info)

import onnx
import numpy as np
from onnx.helper import make_node, make_graph, make_tensor, make_model, make_opsetid

model_path = "onnx_model.onnx"

onnx_model = onnx.load(model_path)

graph = onnx_model.graph
print("graph_input:", graph.input)
print("graph_output:", graph.output)
# print("value_info:", graph.value_info)

tensor_shapes={}
tensor_dtypes={}

value_infos = [val_info for val_info in graph.value_info]
inputs = [input_ for input_ in graph.input]
outputs = [output_ for output_ in graph.output]
value_infos.extend(inputs)
value_infos.extend(outputs)

for val_info in value_infos:
    tensor_name = val_info.name
    dtype = val_info.type.tensor_type.elem_type
    shape = [dim.dim_value for dim in val_info.type.tensor_type.shape.dim]

    tensor_shapes[tensor_name] = shape
    tensor_dtypes[tensor_name] = dtype

print("tensor_shapes:", tensor_shapes)
print("tensor_dtypes:", tensor_dtypes)


def get_tensor_proto_shape(tensor_proto):
    shape = [elem for elem in tensor_proto.dims]
    return shape

图中插入新节点

这里演示了在conv2d_transpose后面插入一个pad算子的过程。

主要流程是创建pad和pad_size的const node并且把这两个node插入到graph合适的位置。

tricks主要是名称的修改，这里有两种方案：第一种是不改变conv2d_transpose的输出名称，插入的新pad输入名称则是conv2d_transpose的原始输出名称，但是pad会引入一个新的output名称，然后这时需要修改原来连接到conv2d_transpose输出的所有node的输入名称。

第二种方案是修改conv2d_transpose的输出名称，而pad算子使用conv2d_transpose原来的输出名称，这样就不用修改conv2d_transpose原来输出node的输入名称。

方案二的优势在于修改node少，并且能够兼容conv2d_transpose输出是graph output的情况，而方案1无法处理。下面代码正是使用了方案2。

注意onnx_model.graph.node.insert(idx, new_node)类似于python list的insert，是把new_node插入到原来位于idx的node的前面，也就是new_node的topo id将为idx。

插入到graph.node的末尾可以直接使用extend。

import onnx
import numpy as np

def create_pad(node, out_idx, pad_size):
    """Create pad and pad size node"""
    node_out_name = node.output[out_idx]
    node_out_new_name = node_out_name + "_padded"
    pad_size_node_name = node_out_name + "_pad_size"
    pad_node_name = node_out_name + "_pad"

    node.output[out_idx] = node_out_new_name

    pad_size = np.array(pad_size).astype("int64")
    pad_size_node = onnx.helper.make_node(
        op_type="Constant",
        inputs=[],
        outputs=[pad_size_node_name],
        name=pad_size_node_name,
        value=onnx.helper.make_tensor(name="const_value",
                                      data_type=onnx.TensorProto.INT64,
                                      dims=pad_size.shape,
                                      vals=pad_size.reshape(-1)))

    pad_node = onnx.helper.make_node(op_type="Pad",
                                     inputs=[node_out_new_name, pad_size_node_name],
                                     outputs=[node_out_name],
                                     name=pad_node_name)

    pas_attr = onnx.helper.make_attribute("mode", 'reflect')
    pad_node.attribute.extend([pas_attr])
    return [pad_size_node, pad_node]

def get_node(onnx_model, node_name):
    for node_id, node in enumerate(onnx_model.graph.node):
        if node.name == node_name:
            return node, node_id
    return None, 0

model_path = "modified.onnx"
onnx_model = onnx.load(model_path)

conv2d_trans, node_id = get_node(onnx_model, "ConvTranspose__9")

pad_size = [0, 0, 0, 0, 0, 0, 1, 0]
new_nodes = create_pad(conv2d_trans, 0, pad_size)

# insert new nodes after the target node with correct topological order
for new_node in reversed(new_nodes):
    onnx_model.graph.node.insert(node_id + 1, new_node)

onnx.checker.check_model(onnx_model)
onnx.save(onnx_model, "test_add_pad.onnx")

删除节点

import onnx

model_path = "bert_model_int32.onnx"
out_model_path = "bert_model_int32_fp32.onnx"

onnx_model = onnx.load(model_path)
graph = onnx_model.graph

# demo for remove node with single input and output
in_rename_map = {}

for node_id, node in enumerate(graph.node):
    if node.name == "Cast_1185":
        in_name = node.input[0]
        out_name = node.output[0]
        in_rename_map = {out_name: in_name}
        del graph.node[node_id]
        break

for node_id, node in enumerate(graph.node):
    for in_id, in_name in enumerate(node.input):
        if in_name in in_rename_map:
            node.input[in_id] = in_rename_map[in_name]

onnx.save(onnx_model, out_model_path)

提取子图

主要工作是原来的graph.node拷贝子图node到新的graph并且创建graph input和output。

Extracting Sub-model with Inputs Outputs Tensor Names

Function extract_model() extracts sub-model from an ONNX model. The sub-model is defined by the names of the input and output tensors exactly.

import onnx

input_path = 'path/to/the/original/model.onnx'
output_path = 'path/to/save/the/extracted/model.onnx'
input_names = ['input_0', 'input_1', 'input_2']
output_names = ['output_0', 'output_1']

onnx.utils.extract_model(input_path, output_path, input_names, output_names)

onnx官方的子图提取工具无法处理>2GB的onnx，这时可以用nvidia的onnx-graphsurgeon工具，其也提供了子图提取功能，可以处理>2GB onnx子图提取：

python3 -m pip install onnx_graphsurgeon --index-url https://pypi.ngc.nvidia.com

https://github.com/NVIDIA/TensorRT/blob/master/tools/onnx-graphsurgeon/examples/03_isolating_a_subgraph/isolate.py

import onnx_graphsurgeon as gs
import numpy as np
import onnx

model = onnx.load("backbone.sim.onnx")
graph = gs.import_onnx(model)

tensors = graph.tensors()

# graph.inputs = [tensors["x1"].to_variable(dtype=np.float32, shape=(1, 3, 224, 224))]
# graph.outputs = [tensors["add_out"].to_variable(dtype=np.float32, shape=(1, 3, 224, 224))]

graph.inputs = [tensors["x1"].to_variable(dtype=np.float32)]
graph.outputs = [tensors["add_out"].to_variable(dtype=np.float32)]

graph.cleanup()

onnx.save(gs.export_onnx(graph), "subgraph.onnx")

修改/替换tensor名称、graph输入输出名称

import onnx

model_path = "where1.onnx"
model_path_out = "where2.onnx"

onnx_model = onnx.load(model_path)
graph = onnx_model.graph

replace_names = {
    "/text_model/Cast_2_output_0": "where_input0",
    "/text_model/Constant_15_output_0": "where_input1",
    "/text_model/Sub_output_0": "where_input2",
    "/text_model/Where_1_output_0": "where_output0",
}

for info in graph.input:
    if info.name in replace_names:
        info.name = replace_names[info.name]

for info in graph.output:
    if info.name in replace_names:
        info.name = replace_names[info.name]

for node in graph.node:
    for idx, _name in enumerate(node.input):
        if _name in replace_names:
            node.input[idx] = replace_names[_name]
    for idx, _name in enumerate(node.output):
        if _name in replace_names:
            node.output[idx] = replace_names[_name]

for tensor in graph.initializer:
    if tensor.name in replace_names:
        tensor.name = replace_names[tensor.name]

onnx.save(onnx_model, model_path_out)

onnx 输入shape信息获取

shape正常是正整数，但是也可能是负数，0和字符串

    for _input in onnx_model.graph.input:
        for i, dim_proto in enumerate(_input.type.tensor_type.shape.dim):
            if dim_proto.HasField("dim_value"):
                pass
            elif dim_proto.HasField("dim_param"):
                pass
            dim_proto.dim_value
            dim_proto.dim_param
        _input.type.tensor_type.elem_type

python - Find input shape from onnx file - Stack Overflow

修改输入shape

import onnx

model_path = "matmul1.onnx"
out_model_path = model_path[:-5] + ".reshape.onnx"

onnx_model = onnx.load(model_path)
graph = onnx_model.graph

print("graph_input:", graph.input)
print("graph_output:", graph.output)

new_shapes = {
    "input": [-1, 2, 3, 4],
}

for _input in graph.input:
    print("_input:", _input.name)
    tensor_shape_proto = _input.type.tensor_type.shape

    new_shape = new_shapes[_input.name]
    # delete old shape
    elem_num = len(tensor_shape_proto.dim)
    for i in reversed(range(elem_num)):
        del tensor_shape_proto.dim[i]

    for i, d in enumerate(new_shape):
        dim = tensor_shape_proto.dim.add()
        if d is None:
            d = -1
        if isinstance(d, int):
            dim.dim_value = d
        elif isinstance(d, str):
            dim.dim_param = d
        else:
            raise ValueError(f"invalid shape: {new_shape}")


print("updated graph_input:", onnx_model.graph.input)
# print("updated graph_output:", onnx_model.graph.output)

onnx.checker.check_model(onnx_model)
onnx.save(onnx_model, out_model_path)

修改输入dtype

import onnx

model_path = "bert_model.onnx"

out_model_path = "bert_model_int32.onnx"

onnx_model = onnx.load(model_path)
graph = onnx_model.graph

print("graph_input:", graph.input)
print("graph_output:", graph.output)

for input in graph.input:
    if input.type.tensor_type.elem_type == onnx.TensorProto.DataType.INT64:
        input.type.tensor_type.elem_type = onnx.TensorProto.DataType.INT32

print("updated graph_input:", onnx_model.graph.input)
print("updated graph_output:", onnx_model.graph.output)

onnx.save(onnx_model, out_model_path)

获取中间结果

方法1

onnxruntime非常扯淡的地方是只能获取onnx_model.graph.output中张量的结果，而不像tensorflow pb模型可以获取任意节点的输出。这给精度调试带来非常多不便。要获取想要的输出，可以把相应张量的输出添加到onnx_model.graph.output中：

def add_graph_outputs(onnx_model, output_names):
    graph = onnx_model.graph
    for out_name in output_names:
        graph.output.append(onnx.ValueInfoProto(name=out_name))
    return onnx_model

也可以给这些output的张量添加shape信息，这里是预先实现的获取shape dtype信息的函数。这需要模型经过了infer shape，并使用上面提到的get_tensor_shapes方法从value_info里面提取tensor shape信息。

import onnx

model_path = "model.onnx"
out_model_path = model_path+".new.onnx"

onnx_model = onnx.load(model_path)
graph = onnx_model.graph

tensor_shapes, tensor_dtypes = get_tensor_shapes(graph)

extract_out_tensor_names = [
    "/backbone/MaxPool_output_0",
]

outputs = [onnx.helper.make_tensor_value_info(
    name=name, elem_type=tensor_dtypes[name], shape=tensor_shapes[name]) for name in extract_out_tensor_names]

graph.output.extend(outputs)

print("outputs:", outputs)

onnx.save(onnx_model, out_model_path)

方法2：onnx新特性，不依赖于onnxruntime：

import numpy as np
from onnx.reference import ReferenceEvaluator

onnx_model = "model.onnx"
onnx_sess = ReferenceEvaluator(onnx_model)
x = np.random.randn(1,3,224,224).astype(np.float32)
out = onnx_sess.run(None, {'input': in_tensor})

试用了这个refence evaluator下bug挺多的。

常量折叠

https://github.com/daquexian/onnx-simplifier

opset version转换

onnx/PythonAPIOverview.md at main · onnx/onnx · GitHub

TensorFlow pb模型修改和优化可以参考另一篇文章：

TensorFlow pb模型修改和优化_Luchang-Li的博客-CSDN博客

pb转onnx

model_path=bert_ner.b2.s128.pb
in_names='input_ids:0,input_mask:0'
out_names=loss/Softmax:0

python -m tf2onnx.convert \
    --input  ${model_path}\
    --output ${model_path}.onnx \
    --inputs ${in_names} \
    --outputs ${out_names} \
    --opset 15