赞
踩
废话不多说线上代码
import os import cv2 import torch import time import threading s_t = time.time() path_weight = '/home/yz5/cq/yolov5_24bit/runs/train/exp2/weights/best.pt' # 如果有多个GPU的话,可以让两个model加载在不同的device上面 model0 = torch.hub.load('/home/yz5/cq/yolov5_24bit', 'custom', path = path_weight, source = 'local', device=0) model1 = torch.hub.load('/home/yz5/cq/yolov5_24bit', 'custom', path = path_weight, source = 'local', device=0) dir_path = 'data/bigImages/cut' img_list = [os.path.join(dir_path, img_name) for img_name in os.listdir(dir_path)] batch_size = 4 batches = [img_list[i:i+batch_size] for i in range(0, len(img_list), batch_size)] def run(model, im, i): # print(i) results = model(im,size = 512) # size是图片的大小 results.print() for i in range(0, len(batches), 2): # 2是进程数 if(i+1<len(batches)): t1 = threading.Thread(target=run, args=[model0, batches[i], i], daemon=True) t2 = threading.Thread(target=run, args=[model0, batches[i + 1], i+1], daemon=True) t1.start() t2.start() t1.join() t2.join() else: t1 = threading.Thread(target=run, args=[model0, batches[i], i], daemon=True) t1.start() t1.join() print("Cost Time:", time.time() - s_t)
简单说一下思想,上述代码其实同时包含了Batch Size和多进程的方式。在一个for循环中使用两个进程加载两个batch,就可以实现多进程的效果了。
具体效果以实际情况为例,比如我的机器是RTX3060+12GB的内存。跑一个detect.pyGPU的利用率就已经爆满了,使用多进程反而增加了总体的检测时间。
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。