onnx基础

初次编辑时间：2024/2/7；最后编辑时间：2024/2/12

定义：ONNX（Open Neural Network Exchange）是一种开放式的文件格式，用于存储训练好的机器学习模型。它使得不同的人工智能框架（如 PyTorch、MXNet、Tensorflow）可以采用相同格式存储模型数据并交互。

---

Basic

当我们加载了一个ONNX之后，我们获得的就是一个ModelProto，它包含了

• 一些版本信息

• 生产者信息

• 一个GraphProto，在GraphProto里面又包含了四个repeated数组，它们分别是

  • node(NodeProto类型)，中存放了模型中所有的计算节点，

  • input(ValueInfoProto类型)，存放了模型的输入节点，

  • output(ValueInfoProto类型)，存放了模型中所有的输出节点，

  • initializer(TensorProto类型)，存放了模型的所有权重参数。

ONNX是把一个网络的每一层或者说一个算子当成节点node，使用这些Node去构建一个Graph，最后将Graph和其它的生产者信息，版本信息等合并在一起生成一个Model，也即是最终的ONNX模型文件。

https://github.com/onnx/onnx/blob/master/onnx/helper.py这个文件非常重要，我们可以利用它提供的make_node，make_graph，make_tensor等等接口完成一个ONNX模型的构建，一个示例如下：

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


# The protobuf definition can be found here:
# https://github.com/onnx/onnx/blob/master/onnx/onnx.proto


# 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', # node name
    ['X', 'pads', 'value'], # inputs
    ['Y'], # outputs
    mode='constant', # attributes
)

# Create the graph (GraphProto)
graph_def = helper.make_graph(
    [node_def],
    'test-model',
    [X, pads, value],
    [Y],
)

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

print('The model is:\n{}'.format(model_def))
onnx.checker.check_model(model_def)
print('The model is checked!')

输出：

The model is:
ir_version: 9
opset_import {
  version: 20
}
producer_name: "onnx-example"
graph {
  node {
    input: "X"
    input: "pads"
    input: "value"
    output: "Y"
    op_type: "Pad"
    attribute {
      name: "mode"
      type: STRING
      s: "constant"
    }
  }
  name: "test-model"
  input {
    name: "X"
    type {
      tensor_type {
        elem_type: 1
        shape {
          dim {
            dim_value: 3
          }
          dim {
            dim_value: 2
          }
        }
      }
    }
  }
  input {
    name: "pads"
    type {
      tensor_type {
        elem_type: 1
        shape {
          dim {
            dim_value: 1
          }
          dim {
            dim_value: 4
          }
        }
      }
    }
  }
  input {
    name: "value"
    type {
      tensor_type {
        elem_type: 1
        shape {
          dim {
            dim_value: 1
          }
        }
      }
    }
  }
  output {
    name: "Y"
    type {
      tensor_type {
        elem_type: 1
        shape {
          dim {
            dim_value: 3
          }
          dim {
            dim_value: 4
          }
        }
      }
    }
  }
}

The model is checked!

这个官方示例为我们演示了如何使用onnx.helper的make_tensor，make_tensor_value_info，make_attribute，make_node，make_graph，make_node等方法来完整构建了一个ONNX模型。需要注意的是在上面的例子中，输入数据是一个一维Tensor，初始维度为[2]，这也是为什么经过维度为[1,4]的Pad操作之后获得的输出Tensor维度为[3,4]。另外由于Pad操作是没有带任何权重信息的，所以当你打印ONNX模型时，ModelProto的GraphProto是没有initializer这个属性的。

---

Huggingface Transformers with ONNX

将Huggingface上的模型转换成onnx格式
optimum-cli export onnx --model distilbert-base-uncased-distilled-squad distilbert_base_uncased_squad_onnx/

将Huggingface上下载的模型转换成onnx格式
optimum-cli export onnx --model opt-125m/ --task text-generation opt-125m_onnx/