单进程多线程调用多GPU的LibTorch推理性能异常问题

1. 用独立CUDA流避免隐式同步阻塞 #

默认情况下，所有CUDA操作都运行在默认流中，默认流会自动同步所有其他流的操作，导致多线程的推理任务被串行执行。为每个GPU分配独立的CUDA流，让推理和拷贝操作在专属流中异步执行，可打破串行阻塞：

import torch
import threading
from queue import Queue

def run_inference(model, input_batch, device, stream, result_queue):
    with torch.cuda.device(device):
        with torch.cuda.stream(stream):
            model.eval()
            with torch.no_grad():
                # 异步将输入移至目标GPU
                input_gpu = input_batch.to(device, non_blocking=True)
                output_gpu = model(input_gpu)
                # 异步拷贝输出到CPU，不阻塞当前流
                output_cpu = output_gpu.to('cpu', non_blocking=True)
                result_queue.put(output_cpu)

# 初始化多GPU环境
device_count = torch.cuda.device_count()
devices = [torch.device(f"cuda:{i}") for i in range(device_count)]
models = [torch.jit.load("your_model.pt").to(dev) for dev in devices]
streams = [torch.cuda.Stream(device=dev) for dev in devices]

# 准备输入批次（按GPU数量拆分）
input_batches = [torch.randn(32, 3, 224, 224) for _ in range(device_count)]
result_queue = Queue()

# 启动多线程推理
threads = []
for idx in range(device_count):
    t = threading.Thread(
        target=run_inference,
        args=(models[idx], input_batches[idx], devices[idx], streams[idx], result_queue)
    )
    threads.append(t)
    t.start()

# 等待所有流执行完成，再处理结果
for stream in streams:
    stream.synchronize()

# 收集结果
results = []
while not result_queue.empty():
    results.append(result_queue.get())

2. 优化CPU拷贝的性能瓶颈 #

to(torch::kCPU)慢的核心原因是同步拷贝+频繁设备切换，可通过以下方式优化：

• 批量拷贝：不要每个推理batch单独拷贝到CPU，攒多个GPU输出后一次性批量拷贝
• 异步拷贝+非阻塞：始终使用non_blocking=True参数，配合独立CUDA流实现异步拷贝，避免等待拷贝完成再执行后续操作
• 延迟拷贝：如果CPU端不需要立即处理输出，可将拷贝操作放到推理任务全部完成后统一执行，减少推理过程中的同步阻塞

3. 验证模型与输入的设备绑定 #

确保每个模型实例和输入数据确实在目标GPU上，避免隐式设备迁移：

• 用print(model.device)验证模型所在设备
• 输入数据提前移至对应GPU，不要在推理前才执行input.to(device)，避免触发隐式同步
• 导出TorchScript时，不要硬编码设备信息（比如不要在模型中写死.to('cuda:0')），保证模型可灵活迁移到任意GPU

4. 切换为多进程方案（线程优化无效时） #

既然多进程调用各GPU的性能符合预期，说明线程模型下的CUDA上下文共享存在同步瓶颈。可采用torch.multiprocessing实现多进程推理：

• 每个子进程加载模型到专属GPU，独立管理CUDA上下文
• 通过队列或管道传递输入输出，避免跨进程的设备同步问题

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