我们在训练好自己的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判断精度是否一致,需要看具体需求.