多线程访问GPU上的相同模型进行推理

时间:2020-05-31 05:23:43

标签: python pytorch

我有一个cnn模型,该模型已加载到GPU上,并且对于每个图像,都必须创建一个新线程并将其分离以在该图像上运行该模型。这可能吗,如果可以,那安全吗?

1 个答案:

答案 0 :(得分:1)

是的,您绝对可以。它有两个方面。如果要并行运行每个模型,则必须在多个GPU中加载相同的模型。如果不需要(只需要线程部分),则可以加载模型并使用concurrent.futures.ThreadPoolExecutor()。在每次通话中,您都可以传递图片。

我用Darknet框架演示了一个例子。

我将模型加载到两个单独的GPU中(对于并行操作,您也可以避免这种情况),并且每次收到请求时,我都使用ThreadPoolExecutor将图像传递给处理函数。

from darknet import *
import concurrent.futures
import time

# you can avoid this part if you don't need multiple GPUs
set_gpu(0) # running on GPU 0
net1 = load_net(b"cfg/yolov3-lp_vehicles.cfg", b"backup/yolov3-lp_vehicles.backup", 0)
meta1 = load_meta(b"data/lp_vehicles.data")

set_gpu(1) # running on GPU 1
net2 = load_net(b"cfg/yolov3-lp_vehicles.cfg", b"backup/yolov3-lp_vehicles.backup", 0)
meta2 = load_meta(b"data/lp_vehicles.data")


def f(x):
    if x[0] == 0: # gpu 0
        return detect_np_lp(net1, meta1, x[1])
    else:
        return detect_np_lp(net2, meta2, x[1])

    

def func2(): # with threading
    a1 = cv2.imread("lp_tester/bug1.jpg")
    a2 = cv2.imread("lp_tester/bug2.jpg")
    nums = [(0, a1), (1, a2)] # the first element in tuple denotes GPU ID
    with concurrent.futures.ThreadPoolExecutor() as executor:
        r_m = [val for val in executor.map(f, nums)]
    print('out f2')
    #return r_m

t1 = time.time()
func2()
t2 = time.time()
print(t2-t1)