python - 多个线程访问 GPU 上的同一模型进行推理

Question

我有一个加载到 GPU 上的 cnn 模型，对于每张图像，必须创建并分离一个新线程以在此图像上运行模型。这可能吗？如果可以，安全吗？

Answer 1

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

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

我将模型加载到两个独立的 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)