基于yolov5的onnx精度测试[附代码]

我们在训练好自己的yolov5模型后,需要对模型进行部署,大多是将torch转为onnx格式进行使用.但在部署之前需要对转的onnx模型进行精度上的测试,看看和torch下的精度是否一致,如果不一致或差距较大还需要进一步的调整.

导出为onnx格式在yolov5/export.py就有实现,输入命令:

python export.py --data data/mydata.yaml --weights [你的torch权重路径]

将会在你的项目中生成对应的onnx模型.可以通过Netron对onnx模型进行可视化.如果你不想下载Netron,可以访问网页版的Netron.

我这里将提供三种方式来测试转化后的onnx精度.

方式1:测试onnx与torch的输出

这种方式最为简单粗暴,可以直接对比onnx和torch的输出,看两者的输出是否一致即可,不过该方法有一定的局限性,测试不够全面,但可以用来做参考.下面附上代码:

import onnxruntime
import torch
from models.experimental import attempt_load
import numpy as np
def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
 
onnx_weights = ''  # onnx权重路径
torch_weights = ''  # torch权重路径
session = onnxruntime.InferenceSession(onnx_weights,providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
model = attempt_load(torch_weights)
 
input1 = torch.randn(1, 3, 640, 640)   # tensor
img = input1.numpy().astype(np.float32)  # array
model.eval()
with torch.no_grad():
    torch_output = model(input1)[0]
# get yolov5 onnx ouput
onnx_output = torch.tensor(session.run([session.get_outputs()[0].name], {session.get_inputs()[0].name: img}))
# 判断输出结果是否一致，小数点后3位一致即可
print(np.testing.assert_almost_equal(to_numpy(torch_output), onnx_output[0], decimal=3))

这里使用随机输入input1进行测试,看torch下的模型和onnx模型输出是否一致.判断方法为采用np.testing.assert_almost_equal进行测试,判断输出的小数点后三位,如果一致,输出结果为None.

方法2:加入后处理的精度测试

方法1的精度测试可以直接看onnx和torch的输出是否一致,但这方法有一定的局限性,因此我们可以用onnx来测试mAP.[用onnx测试mAP在yolov5 6.0中是没有该功能能,因此我们需要修改代码]

在yolo中的后处理其实有很多,比如输出后通过置信度和NMS的滤除就输出后处理.但这里的后处理并不是指NMS的后处理,而是指对三个特征图解码的部分.也就是models/yolo.py中的Detect中的下面这部分对应的代码:

            if not self.training:  # inference  解码，即有后处理
                if self.grid[i].shape[2:4] != x[i].shape[2:4] or self.onnx_dynamic:
                    self.grid[i], self.anchor_grid[i] = self._make_grid(nx, ny, i)
 
                y = x[i].sigmoid()
                if self.inplace:
                    y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i]  # xy
                    y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh
                else:  # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953
                    xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i]  # xy
                    wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh
                    y = torch.cat((xy, wh, y[..., 4:]), -1)
                z.append(y.view(bs, -1, self.no))
        return x if self.training else (torch.cat(z, 1), x)

通过export.py导出onnx模型后修改val.py,我们将用onnx模型来测试mAP.

修改1:加载模型

        # Load model
        if onnx:
            check_requirements(('onnx', 'onnxruntime'))
            import onnxruntime
            session = onnxruntime.InferenceSession(str(weights), None)
      
        gs = 64

修改2:图像预处理

        if onnx:
            img = img.numpy().astype(np.float32)
            names = {k: v for k, v in enumerate(data['names'])}
        else: # torch img:torch.uint8
            img = img.to(device, non_blocking=True)
            img = img.half() if half else img.float()  # uint8 to fp16/32
        img /= 255.0  # 0 - 255 to 0.0 - 1.0  图像的归一化操作
        targets = targets.to(device)  # get标签
        nb, _, height, width = img.shape  # batch size, channels, height, width
        t2 = time_sync()
        dt[0] += t2 - t1  # 图像处理时间

修改3:获得模型的输出

这里的Post_proc是用来判断是否加入后处理的,加入后处理的我将会在方法3中写出来.

        if pt:
            out, train_out = model(img, augment=augment)  # inference and training outputs
        elif onnx:
            if not Post_proc:
                out = torch.tensor(session.run([session.get_outputs()[0].name], {session.get_inputs()[0].name: img}))

这里可以看到session就是我们前面定义好的onnx模型,可以通过get_outputs()和get_inputs()获取输入和输出,因为在yolov5中的输出其实是有两个,一个是通过cat后的拼接输出,一个是列表形式[包含了三个head].

我们需要获得是cat后的输出结果,因此这里是get_outputs()[0],而输入只有1个,也就是在export中定义的名字为'images'的输入.[这个名字翻看export.py中就可以看到]

最后将上面的out送入NMS计算即可:

out = non_max_suppression(out, conf_thres, iou_thres, labels=lb, multi_label=True, agnostic=single_cls)

我们即可获得onnx模型下的mAP结果.然后用该结果和torch下的mAP进行对比.

方法3:不加入后处理的精度测试

方法2是对三个head解码后处理的操作进行mAP的测试,在方法3中是不进行后处理进行精度的测试.可能会想为什么要进行这个操作呢?有方法2不就可以了吗?是这样的,如果你方法2测出的精度和torch是一样,其实大可不比做方法3,但假如你精度不对,那怎么排查呢,我们就可以通过从输出往前一步一步的排查,看看那里不一样,因此我们可以通过方法3看看网络的输出的精度是否一致.

方法也很简单,修改models/yolo.py中的后处理部分代码,如下:

我们只需要获得三个head的sigmoid输出即可.

    def forward(self, x):  # x是list
        z = []  # inference output
        for i in range(self.nl):
            x[i] = self.m[i](x[i])  # conv
            bs, _, ny, nx = x[i].shape  # x(bs,255,20,20) to x(bs,3,20,20,85)
            x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
            if not self.training:
                x[i] = x[i].sigmoid()
        return x
        #     if not self.training:  # inference  解码，即有后处理
        #         if self.grid[i].shape[2:4] != x[i].shape[2:4] or self.onnx_dynamic:
        #             self.grid[i], self.anchor_grid[i] = self._make_grid(nx, ny, i)
        # 
        #         y = x[i].sigmoid()
        #         if self.inplace:
        #             y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i]  # xy
        #             y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh
        #         else:  # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953
        #             xy = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i]  # xy
        #             wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh
        #             y = torch.cat((xy, wh, y[..., 4:]), -1)
        #         z.append(y.view(bs, -1, self.no))
        # return x if self.training else (torch.cat(z, 1), x)

然后在val.py中修改如下代码,获取onnx模型的输出:

        elif onnx:
            if not Post_proc:
                out = torch.tensor(session.run([session.get_outputs()[0].name], {session.get_inputs()[0].name: img}))
            else:
                from tools.make_grid import _make_grid
                # ---------------------------------不加后处理---------------------------------------------------
                out1 = torch.tensor(session.run([session.get_outputs()[0].name], {session.get_inputs()[0].name: img}))
                out2 = torch.tensor(session.run([session.get_outputs()[1].name], {session.get_inputs()[0].name: img}))
                out3 = torch.tensor(session.run([session.get_outputs()[2].name], {session.get_inputs()[0].name: img}))
                out = (out1,out2,out3)
 
                grid = [torch.zeros(1)] * 3  # init grid
                anchor_grid = [torch.zeros(1)] * 3
                stride = torch.tensor([8.,16.,32.])
                z = []
                for i in range(3):
                    _, bs, _, ny, nx, no = out[i].shape
                    if grid[i].shape[2:4] != out[i][0].shape[2:4]:
                        grid[i], anchor_grid[i] = _make_grid(stride, nx, ny, i)
                    y = out[i][0]
                    y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + grid[i]) * stride[i]  # xy
                    y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * anchor_grid[i]  # wh
                    z.append(y.view(bs, -1, no))
                out = torch.cat(z, 1)
                # --------------------------------------------------------------------------------------------
        dt[1] += time_sync() - t2  # 模型推理时间

以上三种精度测试均可以看onnx和torch下的精度是否一致,第一种是看两者输出是否一致[但这个不够全面],后两种是通过mAP判断精度是否一致,需要看具体需求.