使用Python中的onnx.save_model()函数保存生成的模型

在使用Python的ONNX库中，可以使用onnx.save_model()函数保存生成的模型。下面是一个简单的例子，演示了如何使用该函数保存模型。

首先，我们需要安装ONNX库。可以使用以下命令通过pip安装ONNX库：

pip install onnx

接下来，我们将创建一个简单的模型并保存。假设我们要创建一个两层全连接神经网络模型，并将其保存为ONNX格式。

import onnx
import numpy as np
from onnx import helper, TensorProto

# 创建模型
def create_model():
    # 输入和输出节点的名称和形状
    input_name = 'input'
    output_name = 'output'
    input_shape = (1, 1)
    output_shape = (1, 1)

    # 创建输入和输出的张量
    input_tensor = helper.make_tensor_value_info(input_name, TensorProto.FLOAT, input_shape)
    output_tensor = helper.make_tensor_value_info(output_name, TensorProto.FLOAT, output_shape)

    # 创建两个全连接层
    input_data = helper.make_node('Identity', inputs=[input_name], outputs=['input_data'])
    fc_1 = helper.make_node('Gemm', inputs=['input_data', 'weight_1', 'bias_1'], outputs=['fc_1_output'], alpha=1.0, beta=1.0, transB=1)
    relu_1 = helper.make_node('Relu', inputs=['fc_1_output'], outputs=['relu_1_output'])
    fc_2 = helper.make_node('Gemm', inputs=['relu_1_output', 'weight_2', 'bias_2'], outputs=['fc_2_output'], alpha=1.0, beta=1.0, transB=1)
    output = helper.make_node('Identity', inputs=['fc_2_output'], outputs=[output_name])

    # 创建模型的图
    graph_def = helper.make_graph(
        nodes=[input_data, fc_1, relu_1, fc_2, output],
        name='simple_model',
        inputs=[input_tensor],
        outputs=[output_tensor],
        initializer=[
            helper.make_tensor('weight_1', TensorProto.FLOAT, [1, 1]),
            helper.make_tensor('bias_1', TensorProto.FLOAT, [1]),
            helper.make_tensor('weight_2', TensorProto.FLOAT, [1, 1]),
            helper.make_tensor('bias_2', TensorProto.FLOAT, [1])
        ]
    )

    # 创建模型
    model_def = helper.make_model(graph_def, producer_name='onnx-example')

    return model_def

# 保存模型
def save_model(model_def):
    onnx.save_model(model_def, 'simple_model.onnx')
    print("模型已保存为ONNX格式。")

# 创建并保存模型
model = create_model()
save_model(model)

在上面的例子中，我们首先定义了一个create_model()函数，该函数创建了一个简单的两层全连接神经网络模型。然后，我们定义了一个save_model()函数，该函数使用onnx.save_model()函数将模型保存为ONNX格式。最后，我们调用这两个函数来创建并保存模型。

运行上述代码后，会生成一个名为simple_model.onnx的文件，它是我们创建的模型的ONNX格式表示。