.onnx模型示例代码

.onnx模型示例代码通常是用于将已经训练好的模型导出为ONNX格式的。这里以Python中的PyTorch为例：

首先，假设我们有一个已经在PyTorch中训练完成的简单模型（例如一个线性分类器或卷积神经网络）：

import torch
from torch import nn

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 5)  # 假设输入维度为10，输出维度为5

    def forward(self, x):
        return torch.sigmoid(self.fc(x))

# 加载预训练模型
model = SimpleModel()
model.load_state_dict(torch.load('pretrained_model.pth'))

# 获取模型的输入和输出形状
input_shape = (1, 10)  # 假设输入数据形状为(batch_size, input_dim)
output_shape = (1, 5)   # 假设模型输出shape为(batch_size, output_dim)

# 将模型转换为ONNX
torch.onnx.export(model,  # 模型实例
                  torch.randn(input_shape),  # 输入数据
                  'simple_model.onnx',  # 输出文件名
                  export_params=True,  # 导出模型参数
                  opset_version=11,  # ONNX版本
                  do_constant_folding=True,  # 是否折叠常量
                  input_names=['input'],  # 输入名字
                  output_names=['output'],  # 输出名字
                  dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}  # 动态轴定义
)

这段代码首先定义了一个简单的模型，然后加载预训练状态，并使用`torch.onnx.export()`函数将其转换为ONNX格式，保存到'simple_model.onnx'文件中。