单GPU下实现两个PyTorch模型同时推理的技术问询

方法1：利用CUDA流实现异步并行 #

PyTorch的CUDA操作本身支持异步执行，但默认调度下可能还是串行处理。你可以给两个模型分配独立的CUDA流，让GPU同时调度它们的推理任务：

1. 提前把两个模型和输入图像都移到GPU上；
2. 创建两个独立的CUDA流：

   stream1 = torch.cuda.Stream()
   stream2 = torch.cuda.Stream()
   

3. 在各自的流上下文里执行推理：

   # 模型1绑定stream1执行
   with torch.cuda.stream(stream1):
       output1 = model1(input_img)
   # 模型2绑定stream2执行
   with torch.cuda.stream(stream2):
       output2 = model2(input_img)
   # 等待两个流的所有GPU操作完成
   torch.cuda.synchronize()
   

   只要GPU显存和计算资源足够，两个模型的推理会在GPU上并行进行。

方法2：多线程配合CUDA异步触发 #

用Python的threading模块开启两个线程，分别触发两个模型的推理，结合CUDA流的异步特性，让GPU同时接收并处理两个任务：

示例代码：

import threading

def infer_model1(input, result):
    with torch.cuda.stream(torch.cuda.Stream()):
        result['out1'] = model1(input)

def infer_model2(input, result):
    with torch.cuda.stream(torch.cuda.Stream()):
        result['out2'] = model2(input)

# 准备GPU上的输入图像
input_img = input_img.cuda()
result_container = {}

# 启动线程
t1 = threading.Thread(target=infer_model1, args=(input_img, result_container))
t2 = threading.Thread(target=infer_model2, args=(input_img, result_container))
t1.start()
t2.start()

# 等待线程执行结束
t1.join()
t2.join()
# 同步GPU确保结果完全就绪
torch.cuda.synchronize()

output1 = result_container['out1']
output2 = result_container['out2']

这种方式的核心还是依赖CUDA流的异步调度，多线程只是用来同时触发两个模型的推理操作，避免在主线程中串行调用。

关键注意点 #

• 显存限制：两个模型同时加载到GPU，加上输入、输出数据，必须保证显存充足。可以用模型量化（如torch.ao.quantization）压缩模型体积，减少显存占用；
• 任务匹配度：如果两个模型的计算类型差异较大（比如一个是CNN，一个是Transformer），GPU的不同计算单元能更高效地并行利用；若任务类型相似，可能会因为资源竞争导致并行收益有限；
• 实时性优化：摄像头实时输入的预处理步骤也可以放到对应流或线程中异步执行，避免预处理成为整体流程的瓶颈。

内容的提问来源于stack exchange，提问作者soribido