当前位置:   article > 正文

YOLOv5实现多线程推理_yolov5多线程推理

yolov5多线程推理

YOLOv5实现多线程推理

参考链接

官方参考

代码

废话不多说线上代码

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)
  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13
  • 14
  • 15
  • 16
  • 17
  • 18
  • 19
  • 20
  • 21
  • 22
  • 23
  • 24
  • 25
  • 26
  • 27
  • 28
  • 29
  • 30
  • 31
  • 32
  • 33
  • 34
  • 35
  • 36
  • 37
  • 38

思想

简单说一下思想,上述代码其实同时包含了Batch Size和多进程的方式。在一个for循环中使用两个进程加载两个batch,就可以实现多进程的效果了。

具体效果以实际情况为例,比如我的机器是RTX3060+12GB的内存。跑一个detect.pyGPU的利用率就已经爆满了,使用多进程反而增加了总体的检测时间。

声明:本文内容由网友自发贡献,不代表【wpsshop博客】立场,版权归原作者所有,本站不承担相应法律责任。如您发现有侵权的内容,请联系我们。转载请注明出处:https://www.wpsshop.cn/w/小小林熬夜学编程/article/detail/300214?site
推荐阅读
相关标签
  

闽ICP备14008679号