通过onnx.helper构建计算图

onnx是把一个网络的每一层或者说一个算子当成节点node，使用这些node去构建一个graph，即一个网络。通过onnx.helper来生成onnx
步骤如下：
第一步：通过helper.make_tensor_value_info构造出描述输入和输出张量信息的ValueInfoProto对象。要传入张量名、张量的基本数据类型、张量形状这三个信息。
第二步：构造算子节点信息NodeProto，通过在helper.make_node中传入算子类型、输入张量名、输出张量名这三个信息来实现。需要说明的是，ONNX把边的信息保存在了节点信息里，省去了保存边集的步骤。在ONNX中，如果某节点的输入名和之前某节点的输出名相同，就默认这两个节点是相连的。如下面的例子：Mul节点定义了输出c，Add节点定义了输入c，则Mul节点和Add节点是相连的。
第三步：通过helper.make_graph来构造计算图GraphProto。helper.make_graph函数需要传入节点、图名称、输入张量信息、输出张量信息这4个参数。
第四步：通过helper.make_model把计算图GraphProto封装进模型ModelProto里，一个ONNX模型就构造完成了。make_model函数中还可以添加模型制作者、版本等信息。
注释：使用onnx.helper的make_tensor，make_tensor_value_info，make_attribute，make_node，make_graph，make_node等方法来完整构建了一个ONNX模型。

实例一：一个简单的两层全连接网络

import onnx
from onnx import helper
from onnx import TensorProto

def create_model():
    # 使用ONNX helper functions定义两个输入张量'a'和'b'以及一个权重张量'w'
    a = helper.make_tensor_value_info('a', TensorProto.FLOAT, [1,3])
    b = helper.make_tensor_value_info('b', TensorProto.FLOAT, [1,3])
    w1 = helper.make_tensor_value_info('w1', TensorProto.FLOAT, [3,3])
    w2 = helper.make_tensor_value_info('w2', TensorProto.FLOAT, [3,2])

    # 为了使事情简单，我们来创建一些初始权重数据
    # 但是在实际的模型中，应该在训练过程中来更新这些权重值
    weights1 = helper.make_tensor('w1', TensorProto.FLOAT, [3,3], [1, 2, 3, 4, 5, 6, 7, 8, 9])
    weights2 = helper.make_tensor('w2', TensorProto.FLOAT, [3,2], [1, 2, 3, 4, 5, 6])

    fc1 = helper.make_node(
        'Gemm',
        inputs=['a', 'w1', 'b'],
        outputs=['h1'],
        alpha=1.0,
        beta=1.0,
        transB=1
    )

    fc2 = helper.make_node(
        'Gemm',
        inputs=['h1', 'w2', 'b'],
        outputs=['y'],
        alpha=1.0,
        beta=1.0,
        transB=1
    )

    graph = helper.make_graph(
        [fc1, fc2],
        'TwoLayerFC',
        [a, b, w1, w2],
        [helper.make_tensor_value_info('y', TensorProto.FLOAT, [1,2])],
        value_info=[helper.make_tensor_value_info('h1', TensorProto.FLOAT, [1,3])],
        initializer=[weights1, weights2]
    )

    model = helper.make_model(graph, producer_name='onnx-example')
    return model

model = create_model()

# 检查模型并打印
onnx.checker.check_model(model)
print(onnx.helper.printable_graph(model.graph))

实例二：通过conv、relu、add等算子来构建一个简单的模型，结构图如下所示：

代码如下：

import torch
import torch.nn as nn
import onnx
import onnx.helper as helper
import numpy as np

# reference
# https://github.com/shouxieai/learning-cuda-trt/blob/main/tensorrt-basic-1.4-onnx-editor/create-onnx.py

# 构建网络结构
class Model(nn.Module):
    def __init__(self):
        super().__init__()

        self.conv1 = nn.Conv2d(3, 3, 1, 1)
        self.relu1 = nn.ReLU()
        self.conv2 = nn.Conv2d(3, 1, 1, 1)
        self.conv_right = nn.Conv2d(3, 3, 1, 1)
    
    def forward(self, x):
        r = self.conv_right(x)
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.conv2(x + r)
        return x

def hook_forward(fn):

    # @hook_forward("torch.nn.Conv2d.forward")   对torch.nn.Conv2d.forward进行处理
    fnnames   = fn.split(".")     #
    fn_module = eval(".".join(fnnames[:-1]))
    fn_name   = fnnames[-1]
    oldfn = getattr(fn_module, fn_name)
    
    def make_hook(bind_fn):

        ilayer = 0
        def myforward(self, x):
            global all_tensors
            nonlocal ilayer
            y = oldfn(self, x)

            bind_fn(self, ilayer, x, y)
            all_tensors.extend([x, y])   # 避免torch对tensor进行复用
            ilayer += 1
            return y

        setattr(fn_module, fn_name, myforward)
    return make_hook

@hook_forward("torch.nn.Conv2d.forward")
def symbolic_conv2d(self, ilayer, x, y):
    print(f"{type(self)} -> Input {get_obj_idd(x)}, Output {get_obj_idd(y)}")

    inputs = [
        get_obj_idd(x),
        append_initializer(self.weight.data, f"conv{ilayer}.weight"),
        append_initializer(self.bias.data, f"conv{ilayer}.bias")
    ]

    nodes.append(
        helper.make_node(
            "Conv", inputs, [get_obj_idd(y)], f"conv{ilayer}", 
            kernel_shape=self.kernel_size, group=self.groups, pads=[0, 0] + list(self.padding), dilations=self.dilation, strides=self.stride
        )
    )

@hook_forward("torch.nn.ReLU.forward")
def symbolic_relu(self, ilayer, x, y):
    print(f"{type(self)} -> Input {get_obj_idd(x)}, Output {get_obj_idd(y)}")

    nodes.append(
        helper.make_node(
            "Relu", [get_obj_idd(x)], [get_obj_idd(y)], f"relu{ilayer}"
        )
    )

@hook_forward("torch.Tensor.__add__")
def symbolic_add(a, ilayer, b, y):
    print(f"Add -> Input {get_obj_idd(a)} + {get_obj_idd(b)}, Output {get_obj_idd(y)}")

    nodes.append(
        helper.make_node(
            "Add", [get_obj_idd(a), get_obj_idd(b)], [get_obj_idd(y)], f"add{ilayer}"
        )
    )

def append_initializer(value, name):
    initializers.append(
        helper.make_tensor(
            name=name,
            data_type=helper.TensorProto.DataType.FLOAT,
            dims=list(value.shape),
            vals=value.data.numpy().astype(np.float32).tobytes(),
            raw=True
        )
    )
    return name


def get_obj_idd(obj):
    global objmap

    idd = id(obj)
    if idd not in objmap:
        objmap[idd] = str(len(objmap))
    return objmap[idd]

all_tensors = []
objmap = {}
nodes = []
initializers = []

torch.manual_seed(31)
x = torch.full((1, 3, 3, 3), 0.55)
model = Model().eval()
y = model(x)

inputs = [
    helper.make_value_info(
        name="0",
        type_proto=helper.make_tensor_type_proto(
            elem_type=helper.TensorProto.DataType.FLOAT,
            shape=["batch", x.size(1), x.size(2), x.size(3)]
        )
    )
]

outputs = [
    helper.make_value_info(
        name="5",
        type_proto=helper.make_tensor_type_proto(
            elem_type=helper.TensorProto.DataType.FLOAT,
            shape=["batch", y.size(1), y.size(2), y.size(3)]
        )
    )
]

graph = helper.make_graph(
    name="mymodel",
    inputs=inputs,
    outputs=outputs,
    nodes=nodes,
    initializer=initializers
)

# 如果名字不是ai.onnx，netron解析就不是太一样了    区别在可视化的时候，非ai.onnx的名字的话，每一个算子的框框颜色都是一样的
opset = [
    helper.make_operatorsetid("ai.onnx", 11)
]

# producer主要是保持和pytorch一致
model = helper.make_model(graph, opset_imports=opset, producer_name="pytorch", producer_version="1.9")
onnx.save_model(model, "custom.onnx")

print(y)

参考文献：https://github.com/shouxieai/learning-cuda-trt/blob/main/tensorrt-basic-1.4-onnx-editor/create-onnx.py