ref

社区开放麦】第32期 ONNX 新特性和最佳实践介绍 - 知乎

如何修改已有的ONNX模型 - 知乎

ONNX内部节点修改方法_麦克斯韦恶魔的博客-CSDN博客

onnx模型如何增加或者去除里面node，即修改图方法_The space of Shining-CSDN博客

探索发现：tensorflow转onnx时，输入无符号shape的情况解决。_tangshopping的博客-CSDN博客

Creating and Modifying ONNX Model Using ONNX Python API - Lei Mao's Log Book

模型部署入门教程（五）：ONNX 模型的修改与调试 - 知乎

onnx proto, Operators:

https://github.com/onnx/onnx/blob/main/onnx/onnx.proto

https://github.com/onnx/onnx/blob/main/docs/Operators.md

数据类型关系映射

onnx/mapping.py at main · onnx/onnx · GitHub

onnx helper:

onnx/helper.py at main · onnx/onnx · GitHub

大于2GB模型优化方法：

一种大于2GB ONNX模型onnxsim优化方法_Luchang-Li的博客-CSDN博客

onnxsim工具的pass列表：

optimizer/onnxoptimizer/passes at master · onnx/optimizer · GitHub

用法：onnxsim a.onnx b.onnx --skip-optimization fuse_bn_into_conv fuse_pad_into_pool

另一个onnx转换和操作的工具：

GitHub - microsoft/onnxconverter-common: Common utilities for ONNX converters

onnx模型组成：参考onnx.proto里面的各种proto定义，如ModelProto，GraphProto，NodeProto

onnx_model = onnx.load(model_path)
graph = onnx_model.graph
for node_id, node in enumerate(graph.node):
for initializer in graph.initializer:

onnx的常数节点可能是使用initializer表示的，也可能是使用Constant节点表示，内容放到名为"value"的attribute里面(onnx-simplifier能把后者转换为前者，也可以手动进行转换)。

比较怪异的是个别模型部分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

import onnx

onnx_model = onnx.load("models/conv_2d_backprop_static.onnx")

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

for node_id, node in enumerate(graph.node):
    # print node info
    # print(node)
    print("name:", node.name)
    print("op_type:", node.op_type)
    # print("input:", node.input)
    # print("output:", node.output)
    # print("attribute:", node.attribute)

    if node.name == "ConvTranspose__9":
        for attr_id, attr in enumerate(node.attribute):
            # print attr info
            # print("attr:", attr)
            print("attr.name:", attr.name)
            print("attr.type:", attr.type)
            if attr.type == onnx.AttributeProto.AttributeType.INTS:
                print("attr.ints:", attr.ints)
            # if attr.type == onnx.AttributeProto.AttributeType.STRING:
            #     print("attr.s:", attr.s)

            # replace or add attr
            if attr.name == "pads":
                # attr.ints[2] = 0 #  you can also directly modify origin attr
                pas_attr = onnx.helper.make_attribute("pads", [0, 1, 0, 2])
                del node.attribute[attr_id]
                node.attribute.extend([pas_attr])

        print("node new attribute:", node.attribute)

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

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

修改node attr一个更简单方法是删除旧的，添加个新的。
可以给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作为输入输出。

创建一个包含layernorm算子的图：

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新特性：

ONNX Script与基于ONNX Script定义pytorch 导出onnx自定义算子 

https://github.com/microsoft/onnxscript

https://pytorch.org/docs/stable/onnx.html#onnx-script-functions

简化简单模型定义方式

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

2. onnx hub

3. tensorflow pb转onnx模型等

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

onnx_model = shape_inference.infer_shapes(onnx_model)

也可以借助用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)

提取含有少量节点的子图：

import onnx
from onnx.helper import make_node, make_graph, make_tensor, make_model, make_opsetid

# note the model must be infer shape first by onnxruntime shape infer
model_path = "model.onnx"

extract_node_names = [
    "op_name1",
    "op_name2",
]

def get_tensor_shapes(graph):
    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
    return tensor_shapes, tensor_dtypes


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

tensor_shapes, tensor_dtypes = get_tensor_shapes(graph)

init_names = [init.name for init in graph.initializer]
extract_init_names = []
extract_nodes = []


for node_id, node in enumerate(graph.node):
    if node.name in extract_node_names:
        extract_nodes.append(node)
        for in_name in node.input:
            if in_name in init_names:
                extract_init_names.append(in_name)

all_in_names = []
all_out_names = []

for node in extract_nodes:
    for in_name in node.input:
        if in_name in extract_init_names:
            continue
        all_in_names.append(in_name)
    for out_name in node.output:
        all_out_names.append(out_name)

inner_names = set(all_in_names) & set(all_out_names)
valid_in_names = list(set(all_in_names) - inner_names)
valid_out_names = list(set(all_out_names) - inner_names)


input_shapes = [tensor_shapes[in_name] for in_name in valid_in_names]
output_shapes = [tensor_shapes[out_name] for out_name in valid_out_names]
in_dtypes = [tensor_dtypes[in_name] for in_name in valid_in_names]
out_dtypes = [tensor_dtypes[out_name] for out_name in valid_out_names]

initializer = [init for init in graph.initializer if init.name in extract_init_names]

in_num = len(valid_in_names)
out_num = len(valid_out_names)

inputs = [onnx.helper.make_tensor_value_info(
    name=valid_in_names[i], elem_type=in_dtypes[i], shape=input_shapes[i]) for i in range(in_num)]
outputs = [onnx.helper.make_tensor_value_info(
    name=valid_out_names[i], elem_type=out_dtypes[i], shape=output_shapes[i]) for i in range(out_num)]

graph = onnx.helper.make_graph(nodes=extract_nodes, name="sub_graph", inputs=inputs,
                               outputs=outputs, initializer=initializer)

domain = onnx_model.opset_import[0].domain
version = onnx_model.opset_import[0].version

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

print("graph nodes:", len(graph.node))

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

# onnx.checker.check_model(onnx_model)

提取给定输出名称的所有输入节点子图example：

import onnx


def get_input_nodes(graph, node):
    """
    can be accelerated by using a map to store the info
    """
    in_names = node.input
    in_nodes = []
    for node_id, node in enumerate(graph.node):
        for _out_name in node.output:
            if _out_name in in_names:
                in_nodes.append(node)
                break
    return in_nodes


def get_node_cluster(graph, queue):
    visited_nodes = []
    while queue:
        cur_node = queue.pop(0)
        visited_nodes.append(cur_node)

        in_nodes = get_input_nodes(graph, cur_node)
        for node in in_nodes:
            if (node not in queue) and (node not in visited_nodes):
                queue.append(node)
    return visited_nodes


def get_output_nodes(graph, output_names):
    queue = []
    for node_id, node in enumerate(graph.node):
        for _out_name in node.output:
            if _out_name in output_names:
                queue.append(node)
                break
    return queue


model_path = "model.onnx"
out_model_path = "model_extract.onnx"

output_names = [
    "conv1/7x7_s2/bn/sc_2",
]

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

queue = get_output_nodes(graph, output_names)
visited_nodes = get_node_cluster(graph, queue)

sorted_nodes = []
for node in graph.node:
    if node in visited_nodes:
        sorted_nodes.append(node)

all_in_names = []
[all_in_names.extend(node.input) for node in visited_nodes]

initializer_names = [_init.name for _init in graph.initializer]

# get true input without initializer
inputs = []
for _input in graph.input:
    if _input.name in all_in_names:
        inputs.append(_input)

# create output tensor
outputs = [onnx.helper.make_tensor_value_info(
    name=_name, elem_type=onnx.TensorProto.FLOAT, shape=[-1, -1, -1, -1])for _name in output_names]

graph_n = onnx.helper.make_graph(nodes=sorted_nodes, name="extracted_graph", inputs=inputs, outputs=outputs)

# copy initializer
for initializer in graph.initializer:
    if initializer.name in all_in_names:
        graph_n.initializer.append(initializer)

onnx_model_n = onnx.helper.make_model(graph_n)
for item in dir(onnx_model.graph):
    print(item)

# should not use onnx.helper.make_model(graph_n), since the opset is not the same with origin onnx_model
onnx_model.graph.CopyFrom(graph_n)

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

修改/替换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

其他工具