ONNX(Open Neural Network Exchange)介绍

Posted 2022-08-02 fengbingchun

      开放神经网络交换(Open Neural Network Exchange, ONNX)是一种用于表示机器学习模型的开放标准文件格式，可用于存储训练好的模型，它使得不同的机器学习框架(如PyTorch, Caffe等)可以采用相同格式存储模型数据并可交互。ONNX定义了一组和环境、平台均无关的标准格式，来增强各种机器学习模型的可交互性。它让研究人员可以自由地在一个框架中训练模型并在另一个框架中做推理(inference)。

      ONNX是一个社区项目，它以开放源代码的方式托管在GitHub上：https://github.com/onnx/onnx ，最新发布版本为v1.12.0，License为Apache License v2.0，它既可以通过源码编译安装也可以通过简单的pip或conda命令安装。

      通过conda命令安装执行以下命令：

conda install -c conda-forge onnx

      验证是否正确安装执行以下命令：若没有"ModuleNotFoundError: No module named 'onnx'"的错误提示，说明安装成功。

python -c "import onnx"

      ONNX通过Protobuf进行序列化，使用protobuf子集，因此既支持v2也支持v3，参考源码中的onnx.proto或onnx.proto3文件：提供单一文件格式。每个机器学习库都有自己的文件格式，如PyTorch可以保存为'pth'。ONNX为保存和导出模型文件提供了单一标准，该格式文件扩展名为'onnx'。

      (1).定义了一个可扩展的计算图模型(an extensible computation graph model)：各种框架会有不同的图形表示(graph representation)，ONNX为所有这些提供了标准的图形表示。ONNX图通过各种计算节点表示模型图，可以使用Netron等工具进行可视化。

      (2).定义标准数据类型：图中的每个节点通常都有特定的数据类型。为了提供各种框架之间的互操作性，ONNX定义了标准数据类型，包括int、float等。

      (3).定义内置运算符(built-in operators)：这些运算符负责将ONNX中的运算符类型映射到所需的框架。如果要将PyTorch模型转换为ONNX，则所有PyTorch运算符都会映射到它们在ONNX中的关联运算符。

      onnx.proto(或onnx.proto3)中主要的message：

      (1).AttributeProto：包含singular float, integer, string, graph, and tensor values以及repeated float, integer, string, graph, and tensor values的命名属性。

      (2).TensorProto：序列化的张量值。

      (3).SparseTensorProto：序列化的稀疏张量值。

      (4).TensorShapeProto：定义张量shape。维度可以是整数值或符号变量。符号变量代表未知维度。

      (5).ValueInfoProto：定义有关值的信息，包括值的名称、类型和shape。

      (6).NodeProto：计算图由节点(node)的DAG(Directed Acyclic Graph，有向无环图)组成，这些node代表机器学习框架中通常所说的"layer"或"pipeline stage"。例如，node可以是"Conv"类型的节点，它接收image、filter tensor和bias tensor，并产生卷积输出。

      (7).Graphproto：定义了模型的计算逻辑，由参数化的node列表组成，这些node根据其输入和输出形成有向无环图。相当于许多深度学习框架中的"network"或"graph"。

      (8).ModelProto：是一种top-level file/container格式，用于捆绑ML模型并将其计算图与元数据(metadata)相关联。模型的语义由相关的GraphProto描述。

      (9).TrainingInfoProto：用于存储训练模型的信息。定义了两个功能：an initialization-step and a training-algorithm-step

      (10).TypeProto：标准ONNX数据类型。

      ONNX官方教程参考：https://github.com/onnx/tutorials

      在https://github.com/onnx/models中提供了一系列ONNX格式的pre-trained、state-of-the-art模型，包括图像分类、目标检测、图像分割等。

      以下为测试代码：

import numpy as np
import onnx
from onnx import helper, numpy_helper, shape_inference, version_converter
from onnx import AttributeProto, TensorProto, GraphProto

# reference: https://github.com/onnx/onnx/blob/main/docs/PythonAPIOverview.md

def manipulate_tensorproto_and_numpy_array():
    # create a Numpy array
    numpy_array = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=float)
    print("Original Numpy array:\\n\\n".format(numpy_array))

    # convert the Numpy array to a TensorProto
    tensor = numpy_helper.from_array(numpy_array)
    print("TensorProto:\\n".format(tensor))

    # convert the TensorProto to a Numpy array
    new_array = numpy_helper.to_array(tensor)
    print("After round trip, Numpy array:\\n\\n".format(new_array))

    tensorproto_name = "../../data/tensor.pb"
    # save the TensorProto
    with open(tensorproto_name, 'wb') as f:
        f.write(tensor.SerializeToString())
    
    # load a TensorProto
    new_tensor = onnx.TensorProto()
    with open(tensorproto_name, 'rb') as f:
        new_tensor.ParseFromString(f.read())
    print("After saving and loading, new TensorProto:\\n".format(new_tensor))

def run_shape_inference():
    # preprocessing: create a model with two nodes, Y's shape is unknown
    node1 = helper.make_node('Transpose', ['X'], ['Y'], perm=[1, 0, 2])
    node2 = helper.make_node('Transpose', ['Y'], ['Z'], perm=[1, 0, 2])

    graph = helper.make_graph(
        [node1, node2],
        'two-transposes',
        [helper.make_tensor_value_info('X', TensorProto.FLOAT, (2, 3, 4))],
        [helper.make_tensor_value_info('Z', TensorProto.FLOAT, (2, 3, 4))],
    )

    original_model = helper.make_model(graph, producer_name='onnx-examples-2')

    # check the model and print Y's shape information
    onnx.checker.check_model(original_model)
    print("before shape inference, the shape info of Y is:\\n".format(original_model.graph.value_info))

    # apply shape inference on the model
    inferred_model = shape_inference.infer_shapes(original_model)

    # check the model and print Y's shape information
    onnx.checker.check_model(inferred_model)
    print("after shape inference, the shape info of Y is:\\n".format(inferred_model.graph.value_info))

def create_save_onnx_model(model_name):
    # create one input(ValueInfoProto)
    X = helper.make_tensor_value_info('X', TensorProto.FLOAT, [3, 2])
    pads = helper.make_tensor_value_info('pads', TensorProto.FLOAT, [1, 4])
    value = helper.make_tensor_value_info('value', AttributeProto.FLOAT, [1])

    # create one output(ValueInfoProto)
    Y = helper.make_tensor_value_info('Y', TensorProto.FLOAT, [3, 4])

    # create a node(NodeProto): this is based on Pad-11
    node_def = helper.make_node(
        'Pad',                  # name
        ['X', 'pads', 'value'], # inputs
        ['Y'],                  # outputs
        mode='constant',        # attributes
    )

    # create the graph(GraphProto)
    graph_def = helper.make_graph(
        [node_def],        # nodes
        'test-model',      # name
        [X, pads, value],  # inputs
        [Y],               # outputs
    )

    # create the model(ModelProto)
    model_def = helper.make_model(graph_def, producer_name='onnx-example')

    # save the ONNX model
    onnx.save(model_def, model_name)

def load_check_onnx_model(model_name):
    # model is an in-memory ModelProto
    model = onnx.load(model_name)
    print("the model is:\\n".format(model))

    # check the model
    try:
        onnx.checker.check_model(model)
    except onnx.checker.ValidationError as e:
        print("the model is invalid: %s" % e)
    else:
        print("the model is valid!")

def main():
    manipulate_tensorproto_and_numpy_array()
    run_shape_inference()

    model_name = "../../data/example.onnx"
    create_save_onnx_model(model_name)
    load_check_onnx_model(model_name)

    print("test finish")

if __name__ == "__main__":
    main()

      执行结果(部分)如下所示：

      GitHub： https://github.com/fengbingchun/PyTorch_Test