跳转至

13.2 ONNX模型文件

13.2 ONNX模型文件⚓︎

ONNX的全称是开放式神经网络交换(Open Neural Network Exchange),我们已经在前面的快问快答中,对它进行了介绍。既然它是一种开放的模型文件标准格式,那么我们将尝试把前面实现的多入多出三层神经网络保存为ONNX模型文件,以方便在不同的框架中都可以使用。

13.2.1 ONNX模型文件的结构⚓︎

ONNX是一个开放式的规范,定义了可扩展的计算图模型、标准数据类型以及内置的运算符。该文件在存储结构上可以理解为是一种层级的结构,图13-8描述了ONNX模型文件的简化结构。

图13-8 ONNX模型文件的简化结构

最顶层结构是模型(Model),模型记录了该模型文件的基本属性,如使用的ONNX标准的版本、使用的运算符集版本、制造商的名字和版本等信息,除此以外,模型中记录着最主要的信息是图(Graph)。

图(Graph)可以理解为是计算图的一种描述,是由输入、输出以及节点(Node)组成的(图13-8中省略了输入与输出),它们之间通过寻找相同的名字实现连接,也就是说,相同名字的变量会被认为是同一个变量,如果一个节点的输出名字和另一个节点的输入名字相同,这两个节点会被认为是连接在一起的。

节点(Node)就是要调用的运算符,多个节点(Node)以列表的形式在图(Graph)中存储。ONNX支持的运算符类型可以在ONNX官方文档中查看。

13.2.2 创建ONNX节点⚓︎

ONNX采用了Protobuf(Google Protocol Buffer)格式进行存储,这种格式是是一种轻便高效的结构化数据存储格式,可以用于结构化数据序列化,很适合做数据存储或数据交换格式。

Protobuf在使用时需要先在proto文件中定义数据格式,然后构造相关的对象。Python中的onnx库已经提供了创建ONNX的帮助类,使用起来很方便。

下面演示下用make_node API创建一个全连接层节点。注意,由于ONNX的运算符中没有全连接运算符,所以这里用矩阵乘和矩阵加来组成一个全连接层。

首先是矩阵乘,矩阵乘是输入和权重矩阵进行运算,权重矩阵也需要定义一个节点,如下代码用训练好的权重矩阵数据创建了一个常量节点,并添加到节点列表中。其中,对应的输出为fc_weights

weights_node = helper.make_node(
    op_type = "Constant", 
    inputs = [], 
    outputs = ["fc_weights"], 
    value = helper.make_tensor("fc_weights", TensorProto.FLOAT, weights.shape, weights.flatten().astype(float))
  )
node_list.append(weights_node)

然后创建对应的矩阵乘节点,也添加到节点列表中。其中,输入中的fc_weights就是权重矩阵节点的输出,当前矩阵乘节点的输出为matmul_output

matmul_node = helper.make_node(
    op_type = "MatMul",
    inputs = ["input_name", "fc_weights"],
    outputs = ["matmul_output"]
  )
node_list.append(matmul_node)

然后创建矩阵加节点,该节点需要使用偏置矩阵,同样需要先创建偏置矩阵的节点,下面的代码用训练好的偏置矩阵数据创建了一个常量节点,并添加到节点列表中。其中,对应的输出为fc_bias

bias_node = helper.make_node(
    op_type = "Constant", 
    inputs = [], 
    outputs = ["fc_bias"], 
    value=helper.make_tensor("fc_bias", TensorProto.FLOAT, bias.shape, bias.flatten().astype(float))
  )
node_list.append(bias_node)

然后创建对应的矩阵加节点,也添加到节点列表中。其中,输入中的matmul_output是前面矩阵乘节点的输出,fc_bias是偏置矩阵节点的输出,当前矩阵加节点的输出为add_output

matmul_node = helper.make_node(
    op_type = "Add",
    inputs = ["matmul_output", "fc_bias"],
    outputs = ["add_output"]
  )
node_list.append(matmul_node)

这样节点列表中的各节点就创建好了,并且按照输入输出的名字进行了连接,组成了我们需要的全连接层。

13.2.3 创建ONNX文件⚓︎

根据图13-8中的结构,有了节点列表后,就可以创建对应的图(Graph),最终创建出模型(Model),以下就是使用onnx帮助类创建图和模型,并保存在文件中的代码。

graph_proto = helper.make_graph(node_list, "test", input_list, output_list)
model_def = helper.make_model(graph_proto, producer_name="test_onnx")
onnx.save(model_def, output_path)

13.2.4 保存多入多出三层神经网络⚓︎

前面章节中,我们训练了多入多出的三层神经网络来完成MNIST数据集的识别数字任务,现在就可以动手把训练好的网络保存为ONNX格式的模型文件。

这个网络中的三层分别是全连接+Sigmoid、全连接+Tanh、全连接+Softmax,前面已经介绍了如何构造全连接层节点,剩下的就是激活层的构造。其实这些激活函数已经是ONNX内置的运算符,所以对应的节点非常容易构造,对应的代码如下。

if node["type"] in ["Relu", "Softmax", "Sigmoid", "Tanh"]:
  activate_node = helper.make_node(
      node["type"],
      [node["input_name"]],
      [node["output_name"]]
    )
  node_list.append(activate_node)

之前训练时,训练好的权重和偏置分别存储在独立的npz文件中,这里可以重新读取出来,另存为到ONNX文件中。具体的实现这里不再展示,可以到代码目录中获取完整代码并运行,运行之后可以得到mnist.onnx文件,在后面我们将用这个模型文件进行推理。

使用开源工具Netron打开模型文件,可以看到模型结构满足我们的预期。如图13-9所示。

图13-9 多入多出的三层神经网络

代码位置

ch13, Level3