51c~ONNX~合集1

我自己的原文哦~ https://blog.51cto.com/whaosoft/11608027

一、使用Pytorch进行简单的自定义图像分类 ~ONNX 推理

图像分类是计算机视觉中的一项基本任务，涉及训练模型将图像分类为预定义类别。本文中，我们将探讨如何使用 PyTorch 构建一个简单的自定义对象分类模型，然后使用 ONNX 格式将其部署用于推理。

数据集准备

在开始创建模型之前，准备一个标记数据集至关重要。收集要分类的不同对象类别的图像，并根据其类别将它们组织到单独的文件夹中。确保每个类别都有足够数量的图像，以避免过度拟合。

准备如下树所示的数据集

- data- Fruits (dataset name)- train- class 1- class 2- ...- class n- val- class 1- class 2- ...- class n- test- class 1- class 2- ...- class n

我从 kaggle 获取了水果数据集，作为示例链接在此处：

https://www.kaggle.com/datasets/shreyapmaher/fruits-dataset-images

更改 main.py 中“train_dir”和“val_dir”中的路径名

如果需要，初始化数据加载器并添加增强功能。

构建模型

首先导入必要的库，包括 PyTorch。定义自定义对象分类模型的架构，通常使用卷积神经网络 (CNN)。设计网络层，包括卷积层和池化层

import torch.nn as nnclass CustomConvNet(nn.Module):def __init__(self, num_classes):super(CustomConvNet, self).__init__()self.num_classes = num_classesself.layer1 = self.conv_module(3, 16)self.layer2 = self.conv_module(16, 32)self.layer3 = self.conv_module(32, 64)self.layer4 = self.conv_module(64, 128)self.layer5 = self.conv_module(128, 256)self.gap = self.global_avg_pool(256, self.num_classes)def forward(self, x):out = self.layer1(x)out = self.layer2(out)out = self.layer3(out)out = self.layer4(out)out = self.layer5(out)out = self.gap(out)out = out.view(-1, self.num_classes)return outdef conv_module(self, in_num, out_num):return nn.Sequential(nn.Conv2d(in_num, out_num, kernel_size=3, stride=1, padding=1),nn.BatchNorm2d(out_num),nn.LeakyReLU(),nn.MaxPool2d(kernel_size=2, stride=2))def global_avg_pool(self, in_num, out_num):return nn.Sequential(nn.Conv2d(in_num, out_num, kernel_size=3, stride=1, padding=1),nn.BatchNorm2d(out_num),nn.LeakyReLU(),nn.AdaptiveAvgPool2d((1, 1)))

训练模型并进行评估

使用随机梯度下降 (SGD) 或 Adam 优化器等技术实现训练循环，包括前向和后向传播、损失计算和梯度优化。通过跟踪损失和准确度等指标来监控训练过程。我们利用数据增强和正则化等技术来提高模型的泛化能力。

python main.py

for epoch in range(num_epochs):print("Epoch No -", epoch)model.train()running_loss = 0.0running_corrects = 0for inputs, labels in dataLoaders["train"]:# Feeding input and labels to deviceinputs = inputs.to(device, non_blocking=True)labels = labels.to(device, non_blocking=True)optimizer.zero_grad()with torch.set_grad_enabled(True):outputs = model(inputs)_, preds = torch.max(outputs,1)loss = criterion(outputs, labels)loss.backward()optimizer.step()running_loss += loss.item()* inputs.size(0)#calculate accuracyrunning_corrects += torch.sum(preds == labels.data)#scheduler stepexp_lr_scheduler.step()# Calculate average loss and acc for a epochepoch_loss = running_loss/len(train_data)epoch_acc = running_corrects.double()/len(train_data)print('Loss:{} , Acc{}'.format(epoch_loss, epoch_acc))# Saving model every five epochif (epoch%5 == 0):save_model(model,epoch)

运行预测

确保你是否已经更改了权重文件名、训练文件夹、predict.py 中的输入图像

python predict.py

导出为 ONNX 格式

一旦您的模型经过训练并在验证集上表现良好，就可以将其导出为 ONNX 格式进行部署和推理。ONNX（开放神经网络交换）是一种开放标准格式，允许不同深度学习框架之间的互操作性。PyTorch 提供了将模型导出为 ONNX 的工具。

确保你是否更改了export.py中的权重文件名

python export.py

# Now we will save this model.
import torch.onnx
torch.onnx.export(model,img,"./savedModels/custommodel.onnx",export_params=True,opset_version=10,verbose=True,              # Print verbose outputinput_names=['input'],     # Names for input tensoroutput_names=['output'])

使用 ONNX 进行推理

加载已保存的 ONNX 模型并对新的未见过的图像进行推理。使用 ONNX 运行时库加载模型、提供输入数据并获得预测。测量推理时间并将其与 PyTorch 模型的推理时间进行比较，以评估通过 ONNX 优化实现的任何性能改进。

确保你是否已经更改了权重文件名、训练文件夹、predict.py 中的输入图像

python onnx_inference.py

# Load the ONNX model
onnx_model = onnx.load("./savedModels/custommodel.onnx")# Create an ONNX runtime session
ort_session = onnxruntime.InferenceSession("./savedModels/custommodel.onnx")inputs = {"input": trans_img.numpy()}
outputs = ort_session.run(None, inputs)

开发板商城天皓智联

在本教程中，我们探索了使用 PyTorch 构建简单自定义对象分类模型的过程。我们学习了如何训练模型、评估其性能并将其导出为 ONNX 格式以进行推理。通过利用 ONNX 运行时，我们演示了如何高效地对新图像进行推理。有了这些知识，您现在可以将自定义对象分类应用于各种实际应用程序并无缝部署模型。

二、读懂 ONNX、TensorRT、OpenVINO部署框架

本文详细介绍了深度学习模型部署过程中常用的几个框架：ONNX、TensorRT 和 OpenVINO，包括它们的功能、优势以及如何将 PyTorch 模型转换为这些框架支持的格式，旨在提高模型在不同硬件平台上的推理效率和性能。文章还讨论了模型转换过程中可能遇到的问题和相应的解决方案。

这一期主要会分几个点展开：为什么我们做部署的时候要在 torch 上更进一步使用 ONNX，TensorRT，OpenVINO 等部署框架，在做 cv 模型部署的时候。我们怎么部署。在做 LLM 部署的时候，我们又会怎么做呢？

动静转换：Torch上更进一步

Torch

最核心的就是 torch 使用了动态图组网。使用动态组网的好处是。可以使用更偏向 python 语法的格式对模型进行定义。下面就给大家一个常见的网络：

import torch
import torch.nn as nn
import torch.nn.functional as Fclass SimpleNet(nn.Module):def __init__(self, input_size, hidden_size, num_classes):super(SimpleNet, self).__init__()self.fc1 = nn.Linear(input_size, hidden_size)self.fc2 = nn.Linear(hidden_size, hidden_size)self.fc3 = nn.Linear(hidden_size, num_classes)def forward(self, x):x = F.relu(self.fc1(x))x = F.relu(self.fc2(x))x = self.fc3(x)return x# 示例使用
input_size = 784  # 例如，对于MNIST数据集
hidden_size = 128
num_classes = 10model = SimpleNet(input_size, hidden_size, num_classes)
print(model)

大家很容易开心的写出这样代码，但问题是，使用动态图模式。不可避免的带来了一系列问题。对于一个动态图来说，面临了以下三点问题：

性能：

动态图在每次执行时都需要重新构建计算图，这可能导致额外的开销。
静态图只需构建一次，然后可以重复高效执行。

优化难度：

动态图难以进行全局优化，因为图结构在运行时可能会改变。
静态图允许更多的编译时优化，如内存分配优化、算子融合等。

内存使用：

动态图可能需要更多的运行时内存，因为它需要保持 Python 解释器和相关对象的活跃状态。
静态图可以更有效地管理内存，尤其是在推理阶段。

所以从 torch 开始，我们第一步要做的就是动转静。拿到静态图才能更好的做整体性能上的优化！

美好的愿景：ONNX

ONNX，全称 Open Neural Network Exchange，是人工智能领域中一个引人入胜的故事。它的诞生源于一个美好的愿景：在纷繁复杂的深度学习世界中架起一座沟通的桥梁。

2017 年的硅谷，各大科技巨头都在人工智能领域奋力拼搏。Facebook（现在的Meta）和 Microsoft 这两个看似竞争对手的公司，却因为一个共同的梦想走到了一起。他们希望打破AI框架之间的壁垒，让不同平台上训练的模型能够自由迁移。就这样，ONNX 项目应运而生。

听起来 ONNX 是不同模型间完美的桥梁，最后聚合到 ONNX 完成推理是很开心和能接受的事情。但是听起来越完美的事情就面临越多的问题，首先是对 ONNX 来说。ONNX 模型在某些情况下可能比原生框架的模型运行得慢。这主要是因为 ONNX 作为一个中间表示，可能无法充分利用特定硬件或框架的优化特性。想象一下，它就像是一个通用翻译器，虽然能够让不同语言的人交流，但可能会损失一些语言中的微妙之处和效率。

除此之外，AI 领域发展的太快。ONNX 并不一定能很好的表示 torch 中各种各样的算子，导致模型转换成 ONNX 失败。谈回之前简单的网络，我们如何把它转换成 ONNX 形式呢？请看：

# 将模型转换为ONNX格式
import torch.onnx# 创建一个示例输入张量
dummy_input = torch.randn(1, input_size)# 指定ONNX文件的输出路径
output_path = "simple_net.onnx"# 导出模型到ONNX
torch.onnx.export(model,               # 要转换的模型dummy_input,         # 模型的输入样例output_path,         # 输出的ONNX文件路径export_params=True,  # 存储训练好的参数权重opset_versinotallow=11,    # ONNX算子集版本do_constant_folding=True,  # 是否执行常量折叠优化input_names=['input'],   # 输入节点的名称output_names=['output'], # 输出节点的名称dynamic_axes={'input' : {0 : 'batch_size'},    # 批处理维度动态'output' : {0 : 'batch_size'}})print(f"Model has been converted to ONNX and saved as {output_path}")

厂家的秘方：OpenVINO、TensorRT

不同厂家都有自己的推理秘制配方：推理引擎。这种趋势反映了 AI 领域的激烈竞争和快速创新。每家公司都希望在这场技术革命中占据有利地位，而自研推理引擎成为了关键战略。

这种做法的核心原因在于硬件差异化和性能优化。不同公司拥有各自独特的硬件架构，如英特尔的 CPU、NVIDIA 的 GPU 或谷歌的 TPU。为了充分发挥这些硬件的潜力，定制化的推理引擎成为必然选择。这些引擎能够针对特定硬件进行深度优化，实现最佳的性能和效率。这其中，我将为大家简单介绍两种。分别是 OpenVINO 和 TensorRT。

（一）OpenVINO

OpenVINO

让我们先将目光投向 OpenVINO。它的故事始于英特尔的实验室，在那里，一群充满激情的工程师梦想着如何让人工智能的力量触手可及。2018 年，OpenVINO 正式诞生，其名字中的"VINO"代表"Visual Inference and Neural network Optimization"，寓意着它要为视觉智能和神经网络优化开辟一条康庄大道。

OpenVINO 可在英特尔®硬件上扩展计算机视觉和非视觉工作负载，从而最大限度地提高性能。它通过从边缘到云的高性能，人工智能和深度学习推理来加速应用程序。

关于OpenVINO的模型转换

import subprocess
import sysdef convert_onnx_to_openvino(onnx_model_path, output_dir):cmd = [sys.executable,  # 使用当前Python解释器"-m", "mo",  # 调用model optimizer"--input_model", onnx_model_path,"--output_dir", output_dir,"--data_type", "FP32"]subprocess.run(cmd, check=True)print(f"Model has been converted to OpenVINO IR format and saved in {output_dir}")# 使用示例
onnx_model_path = "simple_net.onnx"
output_dir = "openvino_model"convert_onnx_to_openvino(onnx_model_path, output_dir)

这个转换过程和 ONNX 很像，在 OpenVINO 具体执行流程里分为反序列化，输入定义和前向执行几方面。

（二）TensorRT

TensorRT

与此同时，在硅谷的另一端，NVIDIA 的工程师们也在编织着自己的 AI 梦想。2017 年，TensorRT 横空出世，它的名字中的"RT"代表"Runtime"，彰显了它对高性能推理的执着追求。

TensorRT 就像是一位技艺精湛的魔法师，它能够将庞大复杂的神经网络模型变成小巧高效的推理引擎。它的法术可以让模型在 NVIDIA 的 GPU 上飞驰，实现令人瞠目的低延迟和高吞吐量。想象一下，它就像是给AI装上了火箭推进器，让智能决策的速度突破音障。

TensorRT 可用于对超大规模数据中心，嵌入式平台或自动驾驶平台进行推理加速。TensorRT 现已能支持 TensorFlow，Caffe，Mxnet，Pytorch 等几乎所有的深度学习框架，将 TensorRT 和 NVIDIA 的 GPU 结合起来，能在几乎所有的框架中进行快速和高效的部署推理。但可惜，TensorRT 是一个闭源的库。

关于TensorRT模型的转换

我们一般会给 TensorRT 的模型叫为 engine，我们可以使用 trt 提供的命令行工具，trtexec进行转换

trtexec --notallow=simple_net.onnx --saveEngine=simple_net.trt --explicitBatch

推理引擎：类似的执行流程

OpenVINO模型部署分为两个部分：模型优化器和推理引擎。

模型优化器将训练好的模型转换为推理引擎可以识别的中间表达 –IR 文件，并在转换过程中对模型进行优化。推理引擎接受经过模型优化器转换并优化的网络模型，为 Intel 的各种计算设备提供高性能的神经网络推理运算。

TensorRT 模型部署也是分为两个部分：build 和 deployment 。

build：这个阶段主要完成模型转换，将不同框架的模型转换到 TensorRT。模型转换时会完成前述优化过程中的层间融合，精度校准。这一步的输出是一个针对特定 GPU 平台和网络模型的优化过的 TensorRT 模型，这个 TensorRT 模型可以序列化存储到磁盘或内存中。存储到磁盘中的文件称之为 plan file。deployment：将上面一个步骤中的 plan 文件首先反序列化，并创建一个 runtime engine，然后就可以输入数据（比如测试集或数据集之外的图片），然后输出分类向量结果或检测结果。

写在最后

模型部署以加速为最终目的，首先就会抛弃易用性。这里特指静态图，在固定的范围内做极致的优化。除了模型上的优化，不同硬件厂商更会在贴近不同硬件上做各种底层上的优化。以获得在特定芯片上极致的性能。请期待后续部署教程吧～

三、onnxruntime部署YOLOv8分割模型详细教程

本文将详细介绍如何使用onnxruntime框架来部署YOLOv8分割模型，为了方便理解，代码采用Python实现。

0. 引言

我之前写的文章《基于YOLOv8分割模型实现垃圾识别》介绍了如何使用YOLOv8分割模型来实现垃圾识别，主要是介绍如何用自定义的数据集来训练YOLOv8分割模型。那么训练好的模型该如何部署呢？YOLOv8分割模型相比检测模型多了一个实例分割的分支，部署的时候还需要做一些后处理操作才能得到分割结果。

本文将详细介绍如何使用onnxruntime框架来部署YOLOv8分割模型，为了方便理解，代码采用Python实现。

1. 准备工作

「安装onnxruntime」onnxruntime分为GPU版本和CPU版本，均可以通过pip直接安装：

pip install onnxruntime-gpu  #安装GPU版本  pip install onnxruntime  #安装CPU版本

「注意:」 GPU版本和CPU版本建议只选其中一个安装，否则默认会使用CPU版本。

「下载YOLOv8分割模型权重」Ultralytics官方提供了用COCO数据集训练的模型权重，我们可以直接从官方网站https://docs.ultralytics.com/tasks/segment/下载使用，本文使用的模型为yolov8m-seg.pt。

「转换onnx模型」调用下面的命令可以把YOLOv8m-seg.pt模型转换为onnx格式的模型：

yolo task=segment mode=export model=yolov8m-seg.pt format=onnx

转换成功后得到的模型为yolov8m-seg.onnx。

2. 模型部署

2.1 加载onnx模型

首先导入onnxruntime包，然后调用其API加载模型即可:

import onnxruntime as ort  session = ort.InferenceSession("yolov8m-seg.onnx", providers=["CUDAExecutionProvider"])

因为我使用的是GPU版本的onnxruntime，所以providers参数设置的是"CUDAExecutionProvider"；如果是CPU版本，则需设置为"CPUExecutionProvider"。

模型加载成功后，我们可以查看一下模型的输入、输出层的属性：

for input in session.get_inputs():  print("input name: ", input.name)  print("input shape: ", input.shape)  print("input type: ", input.type)  for output in session.get_outputs():  print("output name: ", output.name)  print("output shape: ", output.shape)  print("output type: ", output.type)

结果如下：

input name:  images  
input shape:  [1, 3, 640, 640]  
input type:  tensor(float)  
output name:  output0  
output shape:  [1, 116, 8400]  
output type:  tensor(float)  
output name:  output1  
output shape:  [1, 32, 160, 160]  
output type:  tensor(float)

从上面的打印信息可以知道，模型有一个尺寸为[1, 3, 640, 640]的输入层和两个尺寸分别为[1, 116, 8400]和[1, 32, 160, 160]的输出层。

2.2 数据预处理

数据预处理采用OpenCV和Numpy实现，首先导入这两个包

import cv2  
import numpy as np

用OpenCV读取图片后，把数据按照YOLOv8的要求做预处理

image = cv2.imread("soccer.jpg") image_height, image_width, _ = image.shape input_tensor = prepare_input(image, model_width, model_height) print("input_tensor shape: ", input_tensor.shape)

其中预处理函数prepare_input的实现如下：

def prepare_input(bgr_image, width, height):  image = cv2.cvtColor(bgr_image, cv2.COLOR_BGR2RGB)  image = cv2.resize(image, (width, height)).astype(np.float32)  image = image / 255.0  image = np.transpose(image, (2, 0, 1))  input_tensor = np.expand_dims(image, axis=0)  return input_tensor

处理流程如下：

1. 把OpenCV读取的BGR格式的图片转换为RGB格式；  
2. 把图片resize到模型输入尺寸640x640；  
3. 对像素值除以255做归一化操作；  
4. 把图像数据的通道顺序由HWC调整为CHW；  
5. 扩展数据维度，将数据的维度调整为NCHW。

经过预处理后，输入数据input_tensor的维度变为[1, 3, 640, 640]，与模型的输入尺寸一致。

2.3 模型推理

输入数据准备好以后，就可以送入模型进行推理：

outputs = session.run(None, {session.get_inputs()[0].name: input_tensor})

前面我们打印了模型的输入输出属性，可以知道模型有两个输出分支，其中一个output0是目标检测分支，另一个output1则是实例分割分支，这里打印一下它们的尺寸看一下

#squeeze函数是用于删除shape中为1的维度，对output0做transpose操作是为了方便后续操作  
output0 = np.squeeze(outputs[0]).transpose()  
output1 = np.squeeze(outputs[1])  
print("output0 shape:", output0.shape)  
print("output1 shape:", output1.shape)

结果如下：

output0 shape: (8400, 116)  
output1 shape: (32, 160, 160)

处理后目标检测分支的维度为[8400, 116]，表示模型总共可以检测出8400个目标（大部分是无效的目标），每个目标包含116个参数。刚接触YOLOv8分割模型的时候可能会对116这个数字感到困惑，这里有必要解释一下：每个目标的参数包含4个坐标属性（x,y,w,h）、80个类别置信度和32个实例分割参数，所以总共是116个参数。实例分割分支的维度为[32, 160, 160]，其中第一个维度32与目标检测分支中的32个实例分割参数对应，后面两个维度则由模型输入的宽和高除以4得到，本文所用的模型输入宽和高都是640，所以这两个维度都是160。

2.4 后处理

首先把目标检测分支输出的数据分为两个部分，把实例分割相关的参数从中剥离。

boxes = output0[:, 0:84]  
masks = output0[:, 84:]  
print("boxes shape:", boxes.shape)  
print("masks shape:", masks.shape)

boxes shape: (8400, 84)  
masks shape: (8400, 32)

然后实例分割这部分数据masks要与模型的另外一个分支输出的数据output1做矩阵乘法操作，在这之前要把output1的维度变换为二维。

output1 = output1.reshape(output1.shape[0], -1)  
masks = masks @ output1  
print("masks shape:", masks.shape)

masks shape: (8400, 25600)

做完矩阵乘法后，就得到了8400个目标对应的实例分割掩码数据masks，可以把它与目标检测的结果boxes拼接到一起。

detections = np.hstack([boxes, masks])  
print("detections shape:", detections.shape)

detections shape: (8400, 25684)

到这里读者应该就能理解清楚了，YOLOv8模型总共可以检测出8400个目标，每个目标的参数包含4个坐标属性（x,y,w,h）、80个类别置信度和一个160x160=25600大小的实例分割掩码。

由于YOLOv8模型检测出的8400个目标中有大量的无效目标，所以先要通过置信度过滤去除置信度低于阈值的目标，对于满足置信度满足要求的目标还需要通过非极大值抑制（NMS）操作去除重复的目标。

objects = []  
for row in detections:  prob = row[4:84].max()  if prob < 0.5:  continue  class_id = row[4:84].argmax()  label = COCO_CLASSES[class_id]  xc, yc, w, h = row[:4]  // 把x1, y1, x2, y2的坐标恢复到原始图像坐标  x1 = (xc - w / 2) / model_width * image_width  y1 = (yc - h / 2) / model_height * image_height  x2 = (xc + w / 2) / model_width * image_width  y2 = (yc + h / 2) / model_height * image_height  // 获取实例分割mask  mask = get_mask(row[84:25684], (x1, y1, x2, y2), image_width, image_height)  // 从mask中提取轮廓  polygon = get_polygon(mask, x1, y1)  objects.append([x1, y1, x2, y2, label, prob, polygon, mask])  // NMS  
objects.sort(key=lambda x: x[5], reverse=True)  
results = []  
while len(objects) > 0:  results.append(objects[0])  objects = [object for object in objects if iou(object, objects[0]) < 0.5]

这里重点讲一下获取实例分割掩码的过程。

前面说了每个目标对应的实例分割掩码数据大小为160x160，但是这个尺寸是对应整幅图的掩码。对于单个目标来说，还要从这个160x160的掩码中去截取属于自己的掩码，截取的范围由目标的box决定。上面的代码得到的box是相对于原始图像大小，截取掩码的时候需要把box的坐标转换到相对于160x160的大小，截取完后再把这个掩码的尺寸调整回相对于原始图像大小。截取到box大小的数据后，还需要对数据做sigmoid操作把数值变换到0到1的范围内，也就是求这个box范围内的每个像素属于这个目标的置信度。最后通过阈值操作，置信度大于0.5的像素被当做目标，否则被认为是背景。

具体实现的代码如下：

def get_mask(row, box, img_width, img_height):  mask = row.reshape(160, 160)  x1, y1, x2, y2 = box  // box坐标是相对于原始图像大小，需转换到相对于160*160的大小  mask_x1 = round(x1 / img_width * 160)  mask_y1 = round(y1 / img_height * 160)  mask_x2 = round(x2 / img_width * 160)  mask_y2 = round(y2 / img_height * 160)  mask = mask[mask_y1:mask_y2, mask_x1:mask_x2]  mask = sigmoid(mask)  // 把mask的尺寸调整到相对于原始图像大小  mask = cv2.resize(mask, (round(x2 - x1), round(y2 - y1)))  mask = (mask > 0.5).astype("uint8") * 255  return mask

这里需要注意的是，160x160是相对于模型输入尺寸为640x640来的，如果模型输入是其他尺寸，那么上面的代码需要做相应的调整。

如果需要检测的是下面这个图片：

通过上面的代码可以得到最左边那个人的分割掩码为

但是我们需要的并不是这样一张图片，而是需要用于表示这个目标的轮廓，这可以通过OpenCV的findContours函数来实现。findContours函数返回的是一个用于表示该目标的点集，然后我们可以在原始图像中用fillPoly函数画出该目标的分割结果。

全部目标的检测与分割结果如下：

3. 一点其他的想法

从前面的部署过程可以知道，做后处理的时候需要对实例分割的数据做矩阵乘法、sigmoid激活、维度变换等操作，实际上这些操作也可以在导出模型的时候集成到onnx模型中去，这样就可以简化后处理操作。

首先需要修改ultralytics代码仓库中ultralytics/nn/modules/head.py文件的代码，把Segment类Forward函数最后的代码修改为：

if self.export:  output1 = p.reshape(p.shape[0], p.shape[1], -1)  boxes = x.permute(0, 2, 1)  masks = torch.sigmoid(mc.permute(0, 2, 1) @ output1)  out = torch.cat([boxes, masks], dim=2)  return out  
else:  return (torch.cat([x[0], mc], 1), (x[1], mc, p))

然后修改ultralytics/engine/exporter.py文件中torch.onnx.export的参数，把模型的输出数量改为1个。