上篇博文探究了一下多进程是否能够对YOLOv5模型推理起到加速作用,本篇主要来研究一下如果将图片批量送入网络中进行检测,是否能对网络的推理起到加速作用。
YOLOv5在训练过程中是可以进行分批次训练(batch_size>1),然而在默认的推理过程中,却没有预留batch_size的相关接口,仍然只是单张图一张张进行检测推理。难道批检测推理的速度不会更快吗?下面通过实验来探究。
本文所使用的版本为官方仓库的最新版本(v7.0)。
首先来看看官方源码默认的推理逻辑,在detect.py
文件中,数据集通过LoadImages
实例化一个类。
dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
LoadImages:
class LoadImages:
# YOLOv5 image/video dataloader, i.e. `python detect.py --source image.jpg/vid.mp4`
def __init__(self, path, img_size=640, stride=32, auto=True, transforms=None, vid_stride=1):
if isinstance(path, str) and Path(path).suffix == '.txt': # *.txt file with img/vid/dir on each line
path = Path(path).read_text().rsplit()
files = []
for p in sorted(path) if isinstance(path, (list, tuple)) else [path]:
p = str(Path(p).resolve())
if '*' in p:
files.extend(sorted(glob.glob(p, recursive=True))) # glob
elif os.path.isdir(p):
files.extend(sorted(glob.glob(os.path.join(p, '*.*')))) # dir
elif os.path.isfile(p):
files.append(p) # files
else:
raise FileNotFoundError(f'{p} does not exist')
images = [x for x in files if x.split('.')[-1].lower() in IMG_FORMATS]
videos = [x for x in files if x.split('.')[-1].lower() in VID_FORMATS]
ni, nv = len(images), len(videos)
self.img_size = img_size
self.stride = stride
self.files = images + videos
self.nf = ni + nv # number of files
self.video_flag = [False] * ni + [True] * nv
self.mode = 'image'
self.auto = auto
self.transforms = transforms # optional
self.vid_stride = vid_stride # video frame-rate stride
if any(videos):
self._new_video(videos[0]) # new video
else:
self.cap = None
assert self.nf > 0, f'No images or videos found in {p}. ' \
f'Supported formats are:\nimages: {IMG_FORMATS}\nvideos: {VID_FORMATS}'
def __iter__(self):
self.count = 0
return self
def __next__(self):
if self.count == self.nf:
raise StopIteration
path = self.files[self.count]
if self.video_flag[self.count]:
# Read video
self.mode = 'video'
for _ in range(self.vid_stride):
self.cap.grab()
ret_val, im0 = self.cap.retrieve()
while not ret_val:
self.count += 1
self.cap.release()
if self.count == self.nf: # last video
raise StopIteration
path = self.files[self.count]
self._new_video(path)
ret_val, im0 = self.cap.read()
self.frame += 1
# im0 = self._cv2_rotate(im0) # for use if cv2 autorotation is False
s = f'video {self.count + 1}/{self.nf} ({self.frame}/{self.frames}) {path}: '
else:
# Read image
self.count += 1
im0 = cv2.imread(path) # BGR
assert im0 is not None, f'Image Not Found {path}'
s = f'image {self.count}/{self.nf} {path}: '
if self.transforms:
im = self.transforms(im0) # transforms
else:
im = letterbox(im0, self.img_size, stride=self.stride, auto=self.auto)[0] # padded resize
im = im.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB
im = np.ascontiguousarray(im) # contiguous
return path, im, im0, self.cap, s
该类实例化之后,首先是根据传入的path
类型来判断是单张图片还是文件夹。
然后在循环过程中,执行__next__
方法,此时开始读取文件,并进行Inference、NMS等后续操作。
for path, im, im0s, vid_cap, s in dataset:
with dt[0]:
im = torch.from_numpy(im).to(model.device)
im = im.half() if model.fp16 else im.float() # uint8 to fp16/32
im /= 255 # 0 - 255 to 0.0 - 1.0
if len(im.shape) == 3:
im = im[None] # expand for batch dim
# Inference
with dt[1]:
visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False
pred = model(im, augment=augment, visualize=visualize)
# NMS
with dt[2]:
pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
因此,不难发现,原始的detect.py只使用了单图进行推理。
基本没见到有人做过多图推理的改造探索,在官方仓库的某issue中,找到了作者提供的一种调用思路,稍作改造,代码如下:
import os
import cv2
import torch
import time
s_t = time.time()
# Model
model = torch.hub.load('D:/Desktop/yolov5-master', 'custom', 'yolov5s.pt', source='local')
img_list = []
dir_path = "data/images"
for i in os.listdir(dir_path):
img_list.append(cv2.imread(dir_path + '/' + i)[..., ::-1])
# Inference
results = model(img_list, size=1280) # batch of images
# Results
results.print()
results.save()
print("Cost Time:", time.time() - s_t)
该方法主要调用的是hubconf.py
里面的内容,上面是直接将整个文件夹中的所有图片变成一个batch,输送到网络中进行检测。
实际进行检测的代码块在yolo.py
文件中的_forward_once
方法。
class BaseModel(nn.Module):
def _forward_once(self, x, profile=False, visualize=False):
y, dt = [], [] # outputs
for m in self.model:
if m.f != -1: # if not from previous layer
x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f] # from earlier layers
if profile:
self._profile_one_layer(m, x, dt)
x = m(x) # run
y.append(x if m.i in self.save else None) # save output
if visualize:
feature_visualization(x, m.type, m.i, save_dir=visualize)
return x
这里的x就是输入的Tensor,m是模型的每一层结构,这里不断将输入循环到下一层,实现了网络的批量推理。
下面使用RTX4090单卡进行速度测试,数据集选用VisDrone的部分数据,模型选择YOLOv5s:
测试结果如下表所示:
图片数量 | 直接检测花费时间(s) | 批量检测花费时间(s) |
---|---|---|
100 | 3.767014265060425 | 3.9564051628112793 |
200 | 5.948423385620117 | 6.09602165222168 |
注:我这里的批量检测是直接将所有的图片变成一个batch,200张图片之后,显存基本被占满,更多图像就没有进行测试。
从结果可见,批量检测并没有预期的速度提升,反而比直接单张检测更慢。估计这也是为什么官方不在detect中预留多个batch检测接口的原因。