PyTorch开放神经网络交换(Open Neural Network Exchange)ONNX通用格式模型的熟悉

我们在深度学习中可以发现有很多不同格式的模型文件，比如不同的框架就有各自的文件格式：.model、.h5、.pb、.pkl、.pt、.pth等等，各自有标准就带来互通的不便，所以微软、Meta和亚马逊在内的合作伙伴社区一起搞一个ONNX(Open Neural Network Exchange)文件格式的通用标准，这样就可以使得模型在不同框架之间方便的进行互操作了。
在上节的YOLOv8的目标对象的分类，分割，跟踪和姿态估计的多任务检测实践(Netron模型可视化) 我们在最后有接触到这个格式文件，而且给出了一个 https://netron.app/站点，可以将这个文件的计算图和相关属性都能可视化的给呈现出来。YOLO的模型有自带的方法：

from ultralytics import YOLO
model = YOLO('yolov8n-cls.pt')
model.export(format="onnx")

或者命令行的方式
yolo export model=yolov8n-cls.pt format=onnx

1、PyTorch演示onnx

这里我们来熟悉下在PyTorch中onnx是怎么样的。看一个简单示例：

import torch

class myModel(torch.nn.Module):
    def __init__(self):
        super(myModel, self).__init__()

    def forward(self, x):
        return x.reshape(1,3,64,64)

model = myModel()
x = torch.randn(64,64,3)
torch.onnx.export(model, x, 'newmodel.onnx', input_names=['images'], output_names=['newimages'])

修改输入进来的形状，可以看到通过torch.onnx.export就可以生成onnx格式的模型文件。我们将它上传到上面的站点，将会生成如下的流程图：

我们打印看下export方法有哪些参数：

export(model: 'Union[torch.nn.Module, torch.jit.ScriptModule, torch.jit.ScriptFunction]', args: 'Union[Tuple[Any, ...], torch.Tensor]', f: 'Union[str, io.BytesIO]', export_params: 'bool' = True, verbose: 'bool' = False, training: '_C_onnx.TrainingMode' = <TrainingMode.EVAL: 0>, input_names: 'Optional[Sequence[str]]' = None, output_names: 'Optional[Sequence[str]]' = None, operator_export_type: '_C_onnx.OperatorExportTypes' = <OperatorExportTypes.ONNX: 0>, opset_version: 'Optional[int]' = None, do_constant_folding: 'bool' = True, dynamic_axes: 'Optional[Union[Mapping[str, Mapping[int, str]], Mapping[str, Sequence[int]]]]' = None, keep_initializers_as_inputs: 'Optional[bool]' = None, custom_opsets: 'Optional[Mapping[str, int]]' = None, export_modules_as_functions: 'Union[bool, Collection[Type[torch.nn.Module]]]' = False) -> 'None'

2、加载onnx模型

2.1、安装onnx库

为了能够正确的加载onnx格式文件，需要先安装onnx库，依然推荐加豆瓣镜像安装

pip install onnx -i http://pypi.douban.com/simple/ --trusted-host pypi.douban.com

安装好了之后，我们看下是否成功安装：

import onnx
dir(onnx)
'''
['Any', 'AttributeProto', 'EXPERIMENTAL', 'FunctionProto', 'GraphProto', 'IO', 'IR_VERSION', 'IR_VERSION_2017_10_10', 'IR_VERSION_2017_10_30', 'IR_VERSION_2017_11_3', 'IR_VERSION_2019_1_22', 'IR_VERSION_2019_3_18', 'IR_VERSION_2019_9_19', 'IR_VERSION_2020_5_8', 'IR_VERSION_2021_7_30', 'MapProto', 'ModelProto', 'NodeProto', 'ONNX_ML', 'OperatorProto', 'OperatorSetIdProto', 'OperatorSetProto', 'OperatorStatus', 'Optional', 'OptionalProto', 'STABLE', 'SequenceProto', 'SparseTensorProto', 'StringStringEntryProto', 'TensorAnnotation', 'TensorProto', 'TensorShapeProto', 'TrainingInfoProto', 'TypeProto', 'TypeVar', 'Union', 'ValueInfoProto', 'Version', '_Proto', '__builtins__', '__cached__', '__doc__', '__file__', '__loader__', '__name__', '__package__', '__path__', '__spec__', '__version__', '_deserialize', '_get_file_path', '_load_bytes', '_save_bytes', '_serialize', 'checker', 'compose', 'convert_model_to_external_data', 'defs', 'external_data_helper', 'gen_proto', 'google', 'helper', 'hub', 'load', 'load_external_data_for_model', 'load_from_string', 'load_model', 'load_model_from_string', 'load_tensor', 'load_tensor_from_string', 'mapping', 'numpy_helper', 'onnx_cpp2py_export', 'onnx_data_pb', 'onnx_data_pb2', 'onnx_ml_pb2', 'onnx_operators_ml_pb2', 'onnx_operators_pb', 'onnx_pb', 'os', 'parser', 'printer', 'save', 'save_model', 'save_tensor', 'shape_inference', 'typing', 'utils', 'version', 'version_converter', 'write_external_data_tensors']
'''

2.2、onnx模型信息

正确安装之后，就可以加载模型了。

import onnx

myModel = onnx.load("newmodel.onnx")
#打印整个模型信息
#print(myModel)
output = myModel.graph.output
#打印输出层的信息
print(output)

[name: "newimages"
type {
tensor_type {
elem_type: 1
shape {
dim {
dim_value: 1
}
dim {
dim_value: 3
}
dim {
dim_value: 64
}
dim {
dim_value: 64
}
}
}
}
]

2.3、加载外部数据

需要注意的是，外部数据如果跟模型文件不是在同一个目录里面的话，需要使用load_external_data_for_model来加载外部数据

import onnx
from onnx.external_data_helper import load_external_data_for_model

mymodel = onnx.load('xx/model.onnx', load_external_data=False)
load_external_data_for_model(mymodel, 'datasets/')

这样就可以通过mymodel这个模型来加载指定目录的数据了。

3、保存新模型

我们可以试着在模型的基础上做一些修改，然后保存为一个新的模型文件。

import onnx
from onnx import helper

model = onnx.load('newmodel.onnx')
prob_info = helper.make_tensor_value_info('xx',onnx.TensorProto.FLOAT, [1,3,128,128])
model.graph.output.insert(0, prob_info)
onnx.save(model, 'newmodel2.onnx')

将一个新的节点插入到输出层，我们来看下生成的图：

其中helper里面有很多用法，比如增加节点：

node = helper.make_node("Range",inputs=["start", "limit", "delta"],outputs=["output"])
start = np.float32(1)
limit = np.float32(15)
delta = np.float32(2)
output = np.arange(start, limit, delta, dtype=np.float32)
print(output)#[ 1.  3.  5.  7.  9. 11. 13.]

这里的Range是其中一个操作类型(算子)，用法：Operators 更多的操作类型，有兴趣的可以进去查阅

4、常见对象

在onnix常出现的几个对象，AttributeProto,TensorProto,GraphProto,NodeProto一起来了解下：

4.1、AttributeProto

属性

import onnx
from onnx import helper
from onnx import AttributeProto, TensorProto, GraphProto

arg = helper.make_attribute("this_is_an_int", 1701)
print(arg)
'''
name: "this_is_an_int"
type: INT
i: 1701
'''

当然类型还可以是浮点数、字符串、数组等

4.2、NodeProto

节点

node_proto = helper.make_node("Relu", ["X"], ["Y"])
print(node_proto)
'''
input: "X"
output: "Y"
op_type: "Relu"
'''

其中op_type的算子挺多的，比如上面那个Range的用法。

4.3、AttributeProto和NodeProto结合

看下两者结合使用的一个例子，卷积核为3，步幅为1，填充为1的一个卷积。

node_proto = helper.make_node(
    "Conv", ["X", "W", "B"], ["Y"],
    kernel=3, stride=1, pad=1)
# 属性按顺序打印
node_proto.attribute.sort(key=lambda attr: attr.name)
print(node_proto)
'''
input: "X"
input: "W"
input: "B"
output: "Y"
op_type: "Conv"
attribute {
  name: "kernel"
  type: INT
  i: 3
}
attribute {
  name: "pad"
  type: INT
  i: 1
}
attribute {
  name: "stride"
  type: INT
  i: 1
}
'''

另一种更有可读性的打印：

print(helper.printable_node(node_proto))
%Y = Conv[kernel = 3, pad = 1, stride = 1](%X, %W, %B)

4.4、TensorProto和GraphProto

graph_proto = helper.make_graph(
    [
        helper.make_node("FC", ["X", "W1", "B1"], ["H1"]),
        helper.make_node("Relu", ["H1"], ["R1"]),
        helper.make_node("FC", ["R1", "W2", "B2"], ["Y"]),
    ],
    "MLP",
    [
        helper.make_tensor_value_info("X" , TensorProto.FLOAT, [1]),
        helper.make_tensor_value_info("W1", TensorProto.FLOAT, [1]),
        helper.make_tensor_value_info("B1", TensorProto.FLOAT, [1]),
        helper.make_tensor_value_info("W2", TensorProto.FLOAT, [1]),
        helper.make_tensor_value_info("B2", TensorProto.FLOAT, [1]),
    ],
    [
        helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1]),
    ]
)
print(helper.printable_graph(graph_proto))
'''
graph MLP (
  %X[FLOAT, 1]
  %W1[FLOAT, 1]
  %B1[FLOAT, 1]
  %W2[FLOAT, 1]
  %B2[FLOAT, 1]
) {
  %H1 = FC(%X, %W1, %B1)
  %R1 = Relu(%H1)
  %Y = FC(%R1, %W2, %B2)
  return %Y
'''

5、解析器

5.1、parse_graph

onnx.parser.parse_graph将文本表示，创建成ONNX图形

input = """
   agraph (float[N, 128] X, float[128, 10] W, float[10] B) => (float[N, 10] C)
   {
        T = MatMul(X, W)
        S = Add(T, B)
        C = Softmax(S)
   }
"""
graph = onnx.parser.parse_graph(input)
print(helper.printable_graph(graph))
'''
graph agraph (
  %X[FLOAT, Nx128]
  %W[FLOAT, 128x10]
  %B[FLOAT, 10]
) {
  %T = MatMul(%X, %W)
  %S = Add(%T, %B)
  %C = Softmax(%S)
  return %C
}
'''

5.2、parse_model

onnx.parser.parse_model将文本表示，创建成ONNX模型

input = """
   <
     ir_version: 7,
     opset_import: ["" : 10]
   >
   agraph (float[N, 128] X, float[128, 10] W, float[10] B) => (float[N, 10] C)
   {
      T = MatMul(X, W)
      S = Add(T, B)
      C = Softmax(S)
   }
"""
model = onnx.parser.parse_model(input)
print(model)

'''
ir_version: 7
opset_import {
  domain: ""
  version: 10
}
graph {
  node {
    input: "X"
    input: "W"
    output: "T"
    op_type: "MatMul"
    domain: ""
  }
  node {
    input: "T"
    input: "B"
    output: "S"
    op_type: "Add"
    domain: ""
  }
  node {
    input: "S"
    output: "C"
    op_type: "Softmax"
    domain: ""
  }
  name: "agraph"
  input {
    name: "X"
    type {
      tensor_type {
        elem_type: 1
        shape {
          dim {
            dim_param: "N"
          }
          dim {
            dim_value: 128
          }
        }
      }
    }
  }
  input {
    name: "W"
    type {
      tensor_type {
        elem_type: 1
        shape {
          dim {
            dim_value: 128
          }
          dim {
            dim_value: 10
          }
        }
      }
    }
  }
  input {
    name: "B"
    type {
      tensor_type {
        elem_type: 1
        shape {
          dim {
            dim_value: 10
          }
        }
      }
    }
  }
  output {
    name: "C"
    type {
      tensor_type {
        elem_type: 1
        shape {
          dim {
            dim_param: "N"
          }
          dim {
            dim_value: 10
          }
        }
      }
    }
  }
}
'''

引用来源
github：https://github.com/onnx/onnx