物体检测-系列教程17：YOLOV5 源码解析7 （模型可视化/onnx工具使用/export.py/模型格式转换）

😎😎😎物体检测-系列教程 总目录

有任何问题欢迎在下面留言
本篇文章的代码运行界面均在Pycharm中进行
本篇文章配套的代码资源已经上传
点我下载源码

10、模型可视化

yolov5的工程代码几乎都是用pytorch来编写训练的，生成的网络模型文件也是pt格式（pt即pytorch）

有一个非常好用的工具，叫做netron，可以将预训练的模型文件直接解析成可视化界面，可以非常直观并且详细的查看前向传播过程，netron有一个在线版本（也可以安装本地版本），可以直接导入模型文件进行查看：在线查看网络结构

而netron工具通常比较适合用onnx格式的模型文件格式来进行网络结构可视化的查看，常用的网络模型保存文件类型有 .h5 (HDF5，Keras使用此格式)、.pb (TensorFlow Protocol Buffers，TensorFlow框架)、.pt or .pth (PyTorch TorchScript) 、.onnx (Open Neural Network Exchange，一个开放的标准格式，用于在不同的深度学习框架之间交换模型)、 .params (MXNet)

11、模型文件格式转换export.py

11.1 export.py运行配置参数

在yolov5中，models文件夹有一个脚本叫做export.py，这个是专门用来将pytorch训练出来的.pt格式的模型文件转化为其他格式

这个脚本有3种导出格式，分别是TorchScript 、ONNX 、CoreML ，如果没有出现异常信息，在给出一个pt文件后会生成这3个模型文件

--weights ./models/yolov5s.pt --img 640 --batch 1

这里运行的是export.py，只需要改weights 的参数即可，weights参数指定的是预训练的模型的路径，这里直接使用的简单的5s模型

在运行之前需要安装onnx工具包

pip install onnx

执行脚本成功打印的信息：

Namespace(batch_size=1, device='cpu', dynamic=False, grid=False, img_size=[640, 640], weights='../models/yolov5s.pt')
YOLOv5  2021-4-11 torch 1.8.1+cu111 CPU

Fusing layers... 
Model Summary: 224 layers, 7266973 parameters, 0 gradients, 17.0 GFLOPS

Starting TorchScript export with torch 1.8.1+cu111...
  module._c._create_method_from_trace(
TorchScript export success, saved as ../models/yolov5s.torchscript.pt

Starting ONNX export with onnx 1.15.0...
ONNX export success, saved as ../models/yolov5s.onnx
CoreML export failure: No module named 'coremltools'

Export complete (2.97s). Visualize with https://github.com/lutzroeder/netron.

CoreML 是Apple的框架，我没有安装这个所以没有正确导出

11.2 导包部分

import argparse
import sys
import time
sys.path.append('A:/CV/object_detection/yolo/yolov5/yolov5-5.0/')
import torch
import torch.nn as nn
import models
from models.experimental import attempt_load
from utils.activations import Hardswish, SiLU
from utils.general import set_logging, check_img_size
from utils.torch_utils import select_device
import onnx
import coremltools as ct

这里的第四行代码中，需要指定models模块的路径，因为我运行出错，我直接使用了绝对路径

11.3 命令行参数

parser = argparse.ArgumentParser()
parser.add_argument('--weights', type=str, default='./yolov5s.pt', help='weights path')
parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size')
parser.add_argument('--batch-size', type=int, default=1, help='batch size')
parser.add_argument('--dynamic', action='store_true', help='dynamic ONNX axes')
parser.add_argument('--grid', action='store_true', help='export Detect() layer grid')
parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
opt = parser.parse_args()
opt.img_size *= 2 if len(opt.img_size) == 1 else 1
print(opt)
set_logging()
t = time.time()

1. 创建对象
2. weights，预训练模型
3. img-size，图像尺寸，允许多个值，默认640
4. batch-size，默认1
5. dynamic，onnx动态轴
6. grid，在导出模型时包含Detect层的网格
7. device，运行设备，cpu即可
8. opt，解析参数
9. 如果–img-size参数只指定了一个值，则将该值扩大2倍，用于支持只指定一个尺寸时的图像大小调整
10. 打印解析后的命令行参数
11. 设置日志记录的配置
12. 记录当前时间

11.4 模型准备

这段主要加载v5训练好的模型，并对模型输入进行准备

device = select_device(opt.device)
model = attempt_load(opt.weights, map_location=device)
labels = model.names
gs = int(max(model.stride))
opt.img_size = [check_img_size(x, gs) for x in opt.img_size]
img = torch.zeros(opt.batch_size, 3, *opt.img_size).to(device)

1. device ，运行设备
2. model ，使用attempt_load函数加载预训练模型
3. labels ，标签名称列表
4. gs，最大卷积步长
5. 调用check_img_size函数检查并调整图像尺寸，确保每个图像尺寸都是步长的倍数
6. img，创建全为0的tensor作为模型输入

11.5 模型更新

for k, m in model.named_modules():
    m._non_persistent_buffers_set = set()
    if isinstance(m, models.common.Conv):
        if isinstance(m.act, nn.Hardswish):
            m.act = Hardswish()
        elif isinstance(m.act, nn.SiLU):
            m.act = SiLU()
model.model[-1].export = not opt.grid
y = model(img)

1. 遍历模型的所有模块和其名称，通过torch.nn的named_modules方法来返回名称k和模块本身m
2. 设置空集解决版本间的兼容性问题
3. 是否是卷积模块：
4. 如果卷积的激活函数是Hardswish：
5. 替换为友好的Hardswish实现，“友好的实现”（export-friendly implementation）在深度学习模型导出和部署的语境中，通常指的是一种被优化以便更容易地转换或兼容不同深度学习框架（如PyTorch到ONNX，再到CoreML等）的实现方式。对于激活函数如Hardswish和SiLU（Swish），"友好的实现"意味着这些函数的实现方式被调整或重写，以确保它们在模型转换过程中能够被正确识别和支持，尤其是在不同框架间转换时
6. 如果卷积的激活函数是SiLU：
7. 替换为友好的SiLU实现
8. 设置模型的Detect层的导出模式
9. 对模型进行一次“干运行”，即在不进行梯度更新的情况下通过模型前向传播一个零张量。这通常用于确保模型修改没有引入错误，并且模型能够正常处理输入

通过替换特定的激活函数并设置Detect层的导出模式，为YOLOv5模型的导出做准备。这些步骤确保了模型能够与不同版本的PyTorch兼容，同时也使得模型能够更容易地被转换为其他格式，如ONNX或CoreML，以适应不同的部署需求

11.6 TorchScript 格式导出

try:
    print('\nStarting TorchScript export with torch %s...' % torch.__version__)
    f = opt.weights.replace('.pt', '.torchscript.pt')
    ts = torch.jit.trace(model, img)
    ts.save(f)
    print('TorchScript export success, saved as %s' % f)
except Exception as e:
    print('TorchScript export failure: %s' % e)

1. try
2. 打印开始导出TorchScript格式模型的信息，包括当前使用的PyTorch版本
3. f，定义导出的TorchScript模型文件名，将原始模型文件名中的后缀.pt替换为.torchscript.pt
4. ts，使用torch.jit.trace函数对模型进行追踪，生成一个TorchScript对象。这一步需要提供模型(model)和一组输入数据(img)，通过模型的前向传播来追踪其计算图
5. 将追踪得到的TorchScript对象保存到文件系统中
6. 打印保存信息
7. 任何导致导出失败的异常
8. 打印异常信息

11.7 ONNX 格式导出

try: 
    print('\nStarting ONNX export with onnx %s...' % onnx.__version__)
    f = opt.weights.replace('.pt', '.onnx')
    torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'],
                      output_names=['classes', 'boxes'] if y is None else ['output'],
                      dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'},  
                                    'output': {0: 'batch', 2: 'y', 3: 'x'}} if opt.dynamic else None)
    onnx_model = onnx.load(f)
    onnx.checker.check_model(onnx_model)
    print('ONNX export success, saved as %s' % f)
except Exception as e:
    print('ONNX export failure: %s' % e)

1. 打印导入ONNX格式模型的信息
2. f，换模型文件名后缀
3. 使用torch.onnx.export函数将PyTorch模型转换为ONNX格式。这包括指定模型(model)、输入数据(img)、输出文件名(f)、操作集版本(opset_version=12)、输入输出名、以及（可选的）动态轴配置。动态轴允许模型在不同的批次大小和图像长宽下运行
4. onnx_model ，加载导出的ONNX模型文件，以便进行进一步的检查
5. 使用ONNX的检查器验证模型的结构和数据完整性，确保导出的模型没有问题
6. 打印导出成功和保存位置的信息
7. 出现任何导出异常
8. 打印导出失败的异常信息