Ricky Yel Blog

使用onnx库创建一个单Conv节点的模型

 
  • Ricky Yel
  • 2024年 05月08日

通过创建一个单Conv节点的onnx模型来熟悉onnx库的使用和相关的概念。

make_graph 原型

def make_graph(
    nodes: Sequence[NodeProto],
    name: str,
    inputs: Sequence[ValueInfoProto],
    outputs: Sequence[ValueInfoProto],
    initializer: Optional[Sequence[TensorProto]] = None,
    doc_string: Optional[str] = None,
    value_info: Optional[Sequence[ValueInfoProto]] = None,
    sparse_initializer: Optional[Sequence[SparseTensorProto]] = None,
) -> GraphProto:
    """Construct a GraphProto

    Args:
        nodes: list of NodeProto
        name (string): graph name
        inputs: list of ValueInfoProto
        outputs: list of ValueInfoProto
        initializer: list of TensorProto
        doc_string (string): graph documentation
        value_info: list of ValueInfoProto
        sparse_initializer: list of SparseTensorProto
    Returns:
        GraphProto
    """

从函数原型中，可以看出一个onnx模型的网络拓扑图需要包含以下几种类型的数据：

  1. nodes: Sequence[NodeProto]，即节点列表。
  2. name: str，即模型的名称。
  3. inputs: Sequence[ValueInfoProto]，即输入节点列表，主要包含输入节点的名称，形状和数据类型信息。
  4. outputs: Sequence[ValueInfoProto]，即输出节点列表，主要包含输出节点的名称，形状和数据类型信息。
  5. initializer: Optional[Sequence[TensorProto]]，有些类型的节点除了变化的input tensor之外，还有一些固定的输入，这些固定的输入就是initializer。
  6. doc_string: Optional[str]，模型的文档字符串（不甚了解）。
  7. value_info: Optional[Sequence[ValueInfoProto]]，节点间流动的tensor的信息，和输入输出类似，主要包括tensor的名称，形状和数据类型。
  8. sparse_initializer: Optional[Sequence[SparseTensorProto]]，稀疏张量的初始化信息（不甚了解）。
所以说想要创建一个网络，就是准备好上述数据即可

创建一个单Conv节点的模型

import torch
import numpy as np
import onnx
from onnx import helper
from onnx import AttributeProto, TensorProto, GraphProto, ValueInfoProto

# 创建i/o ValueInfoProto
input_name = 'input'
input_tensor = helper.make_tensor_value_info(
    input_name,  # 输入的名称
    TensorProto.FLOAT,  # 数据类型
    [1, 3, 224, 224]  # 数据的形状
)
output_name = 'output'
output_tensor = helper.make_tensor_value_info(
    output_name,  # 输出的名称
    TensorProto.FLOAT,  # 数据类型
    [1, 12, 112, 112]  # 数据的形状
)

# 从现有的权重张量创建初始化器节点
weight_name = 'weight'
weight_data = np.random.randn(12, 3, 2, 2).astype(np.float32)
weight_initializer = onnx.numpy_helper.from_array(weight_data, name=weight_name)

# 创建全零偏置的初始化器节点
bias_name = 'bias'
bias_data = np.zeros((12,), dtype=np.float32)
bias_initializer = onnx.numpy_helper.from_array(bias_data, name=bias_name)

# 创建(12,3,2,2)卷积节点的属性
kernel_shape = [2, 2]  # 卷积核的形状
strides = [2, 2]  # 步幅
pads = [0, 0, 0, 0]  # 填充
dilations = [1, 1]  # 膨胀率
group = 1  # 常规卷积，不分组
auto_pad = 'NOTSET'  # 自动填充方式，默认为NOTSET

# 创建卷积节点
conv_node = helper.make_node(
    'Conv',  # 节点类型
    inputs=[input_name, weight_name, bias_name],  # 输入的名称，需要根据实际情况修改
    outputs=[output_name],  # 输出的名称，需要根据实际情况修改
    name='conv',  # 节点名称
    kernel_shape=kernel_shape,
    strides=strides,
    pads=pads,
    dilations=dilations,
    group=group,
    auto_pad=auto_pad
)

graph = helper.make_graph(nodes=[conv_node],  # 节点
                          name='conv',  # 图的名称
                          inputs=[input_tensor],  # 输入
                          outputs=[output_tensor],  # 输出
                          initializer=[weight_initializer, bias_initializer]  # 初始化器
                          )

# 创建ONNX模型
model = helper.make_model(graph, producer_name='onnx-example')

onnx.checker.check_model(model)

# 保存ONNX模型
onnx.save_model(model, 'conv_model.onnx')

ps: Conv节点的输入输出只需要指定 ValueInfoProto 或 TensorProto 的名称即可，不需要指定具体的数据。

最终拓扑结构为：

            input
              |
              |  weight bias
              |   |     |
             (----Conv---)
                  |
                output

可以通过netron可视化保存的onnx模型文件查看模型的拓扑结构。

搜索