Yolov5 将图片合成batch进行推理-2

Yolov5 将图片合成batch进行推理-2

之前的一篇博客Yolov5 将图片合成batch进行推理-1介绍了一个简单的用法。

这一篇主要就是将路径下的图片以batchsize=n这样送去给model进行检测。

我之前是对下面这种实现方式有误解。下面这个代码其实就只加载了一次模型，和detect.py中多次运行run是不一样的。那是因为在detect.py中的run中有导入模型的步骤。

下面使用我自己训练出来的YOLOv5x进行测试的，路径下一共有290张图片。

import os
import cv2
import torch
import time

s_t = time.time()

# 记载模型
path_weight = '/home/yz5/cq/yolov5_24bit/runs/train/exp3/weights/best.pt'
model = torch.hub.load('/home/yz5/cq/yolov5_24bit', 'custom', path = path_weight, source = 'local', autoshape = True)

# 将图片分开
dir_path = 'data/PCBDatasets/images'  # 你的图片所在的目录
img_list = [os.path.join(dir_path, img_name) for img_name in os.listdir(dir_path)]  # 获取所有图片的路径
batch_size = 1
batches = [img_list[i:i+batch_size] for i in range(0, len(img_list), batch_size)]  # 将图片分成多个batch

for batch in batches:
    results = model(batch, size=512)  # batch of images
    results.print()

print("Cost Time：", time.time() - s_t)

batchsize	time
1	21.7
8	21.1
16	20.8
24	20.8
32	21.1

通过上面那个实验表格其实就可以知道，效果有但是不多。最重要的是，通过使用官方的detecct.py测试这290张图只消耗了15s的时间。可以说这种办法基本上不值得倡导。

如果要深究n个batchsize的方式去做检测的话，那就只能在detect.py上进行修改了。