YOLOv3 从入门到部署（四）YOLOv3模型导出onnx（基于pytorch）

2023-10-31

YOLOv3 从入门到部署（四）YOLOv3模型导出onnx（基于pytorch）

文章目录

YOLOv3 从入门到部署（四）YOLOv3模型导出onnx（基于pytorch）

概述

本篇博客我们重点讲解如何将“YOLOv3 从入门到部署（二）”中构建的模型转换为onnx。关于onnx的介绍，读者可以查询其他资料。本章节讲述的内容需要大量参考YOLOv3 从入门到部署（二）。

DNN模块是opencv的一个深度学习推理模块，我们使用的是opencv 4.5.1。我们最终要使用DNN模块加载我们导出的onnx。要完成这件事我们必须要确保两件事：onnx的导出过程是成功的；onnx可以被DNN正确加载。为什么这么说呢？因为DNN模块尚不完整，有很多onnx中的网络层DNN是不支持的，这会导致即使onnx导出成功，也无法被DNN正确加载。

pytorch导出onnx采坑

pytorch导出onnx有很多资料，但是我也相信很多人在将自己的模型导出为onnx的时候总是会遇到各种各样的问题。我把自己的一些经验和解决办法总结如下（有些经验不只是为了成功导出onnx，也是为了让DNN正确加载onnx）：

尽量不要出现tensor转为numpy或者python的数据类型：打个比方，pytorch的主角是tensor，在转为onnx的时候，pytorch会根据tensor的流向去构建onnx。一旦tensor的流向中断，onnx就会出现警告或者报错。那么什么时候tensor的流向会中断呢？就是在进行数据类型转换的时候。因此不要把tensor转换为numpy或者python的int，float等类型，更不要把tensor转换成列表。
tensor不要和python的数据类型直接做运算：比如我们的程序中有一句self.stride = torch.floor_divide(self.img_dim, self.grid_size)，正常的我们可能会直接使用self.img_dim/self.grid_size然后取整，但是如果要转为onnx，这么写就会报出warning。这是因为在我们的程序中，self.img_dim/是python的int类型，而self.grid_size=inputs.size(2)是一个tensor。这样的话再运算过程中，tensor会被隐式地转换为python的数据类型，转为onnx的时候会被警告有tensor被转换为其它数据类型。
不要使用tensor的广播机制，因为DNN不支持：（这一点主要是为了让DNN正确加载onnx，不遵守这一条不会导致转onnx出问题。）4.5.1的版本下，DNN是不支持广播机制。作为广播机制，就是说在tensor和tensor做运算的时候，两者维度可以不一样。pytorch会根据广播机制自动将缺失的维度补全。但是DNN不吃这一套，所以要想办法尽量避免使用广播机制。比如https://github.com/eriklindernoren/PyTorch-YOLOv3中源程序是

X = x.data + self.grid_x

其中x是一个有着4个维度的tensor，而self.grid_x只是一个二维的tensor，在做加法的时候pytorch会默认使用广播机制，给x的每个元素加上self.grid_x，但是这样的操作在DNN中不被支持。因此我们将其改为

FloatTensor = torch.cuda.FloatTensor if inputs.is_cuda else torch.FloatTensor
X = FloatTensor()           # x 和 self.grid_x维度并不完全相同，为了转onnx成功，需要写成这样
     for i in range(self.num_anchors):
         X = torch.cat( (X, torch.add(x[:, i:i+1, :, :], self.grid_x)), 1)

我们使用for循环逐次和x的元素相加；并且使用torch.add()，而不是普通的加号。

DNN不支持torch.arange()，torch.squeeze()，torch.exp()：目前4.5.1的opencv的DNN并不支持这些操作，甚至更多操作。但是以后的opencv可能会逐渐支持。解决的办法是能避免使用就避免使用，实在没办法，就在后续使用DNN加载onnx前，使用DNN相关函数自定义这些不被支持的模块。在我们的程序中，由于我们无论如何都要使用exp，因此我们决定在DNN中自定义exp模块，具体方法后边展示。源代码中使用

self.grid_x = torch.arange(g).repeat(g, 1).view([1, 1, g, g]).type(FloatTensor)

创建self.grid_x；为了避免使用arange，我们将代码改为如下：

self.grid_x = FloatTensor([i for j in range(self.grid_size) for i in range(self.grid_size)])\
            .view([1, 1, self.grid_size, self.grid_size])

不要对切片进行赋值操作：onnx不支持对切片进行赋值操作，因此要避免使用。
onnx支持多输入多输出：onnx支持模型有多个输入和输出。在我们的模型中，输入只有一个，但是输出有两个。尽管我们将两个yolo层的输出结果保存在一个列表中，但是模型输出的tensor是有两个，因此相当于两个输出。

转onnx代码

conifg_path = "./configs/yolo-fastest-xl.cfg"
weights_path = "./weights/yolo-fastest-xl.weights"
save_path = "./weights/yolo-fastest-xl.onnx"


net = YOLOv3(conifg_path)
# If specified we start from checkpoint
if weights_path:
    if weights_path.endswith(".pth"):
        net.load_state_dict(torch.load(weights_path))       # 加载pytorch格式的权重文件
        print("load_state_dict")
    else:
        net.load_darknet_weights(weights_path)          # 加载darknet格式的权重文件。以.weight为后缀
        print("load_darknet_weights")

net.eval()
inputs = torch.rand(1, 3, 320, 320)
torch.onnx.export(net, inputs, save_path, input_names=["input"], output_names=["outputs0", "outputs1"],
                  verbose=True, opset_version=11)
model = onnx.load(save_path)
onnx.checker.check_model(model)

torch.onnx.export：input_names和output_names要对应模型输出输出的额个数，具体的名字可以自己随便起；verbose=True会在转onnx的过程中打印转换细节；opset_version=11我目前也不太了解。
inputs = torch.rand(1, 3, 320, 320)：这里的输入图像是一张320*320的三通道图片，与最终模型推理时的输入尺寸要一致。
onnx.checker.check_model(model)：验证onnx。

使用DNN加载onnx进行验证

我们前面说过需要自行定义DNN中的exp模块，具体方法如下

import
# 添加 ExpLayer
class ExpLayer(object):
    def __init__(self, params, blobs):
        super(ExpLayer, self).__init__()

    def getMemoryShapes(self, inputs):
        return inputs

    def forward(self, inputs):
        return [np.exp(inputs[0])]
cv2.dnn_registerLayer('Exp', ExpLayer)

# opencv dnn加载
net = cv2.dnn.readNetFromONNX(save_path)
img = inputs.numpy() * 255
img = img[0]
img = img.transpose((1, 2, 0))
img = img.astype('uint8')
blob = cv2.dnn.blobFromImage(img, size=(320, 320))      # img 必须是uint8
print(blob.shape)
net.setInput(blob)
out_blob = net.forward(net.getUnconnectedOutLayersNames())
print(out_blob[1].shape)

out = cv2.dnn.imagesFromBlob(out_blob[1])
print(out[0].shape)

cv2.dnn.blobFromImage：这个函数相当于一个预处理函数，可以对input进行归一化，resize等操作。
out_blob是一个列表，有两个元素，分别是两个yolo层的输出
cv2.dnn.imagesFromBlob：这里需要仔细讲一下这个函数。在C++的opencv中，图像实用Mat类型来存放的，Mat是一个二维矩阵，通道数是channel，因此Mat可以看做是一个有着三个维度的矩阵。但是我们模型的输入是四个维度的矩阵，多出来的维度表示样本数量；因此我们才会使用blobFromImage将一个普通的3*320*320的图片转换为blob，维度为1*3*320*320；在C++中就是将三个维度的Mat转换为四个维度的blob（不过注意blob也是Mat类型）。模型的输出也是四个维度的blob，而yolo层的输出其实是一个二维矩阵，行数表示目标数量。列数是85。因此我们需要使用imagesFromBlob，再将四个维度的blob转换为三个维度的Mat。在C++中，如果一个blob的维度是4*3*320*320，那么imagesFromBlob会返回一个vector，其中有四个元素，每个元素是一个维度为3*320*320的Mat。在后续的博客中，我们使用C++ opencv部署onnx的时候会继续讲这件事。

本文内容由网友自发贡献，版权归原作者所有，本站不承担相应法律责任。如您发现有涉嫌抄袭侵权的内容，请联系:hwhale#tublm.com(使用前将#替换为@)