使用 PyTorch 进行两流 CNN 并行推理

Question

目前我正在完成我的最后一年项目，该项目涉及开发一个多流 CNN 来执行动作识别。但是，最终输出依赖于独立流（空间和时间）生成的输出。我的目标是使推理过程尽可能高效，因此我希望使 2 个不同的流同时运行。默认情况下，它会顺序运行forward函数，因此执行时间会很长。

rgb = network1(input1)
of = network2(input2)
final_output = (rgb + of)/2
return final_output

我浏览了有关 PyTorch 多处理的一些信息，并尝试了一些使用 torch.multiprocessing.Process 的示例，但是执行时间似乎比我预期的要长。代码如下所示。

import torch
import torchvision
import torch.multiprocessing as mp
import time

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

net1 = torchvision.models.quantization.mobilenet_v3_large(pretrained=True,quantize=False)
net2 = torchvision.models.quantization.mobilenet_v3_large(pretrained=True,quantize=False)

if __name__ == "__main__":
    inputs = torch.rand(1, 3, 224, 224)
    start = time.time()
    outputs = net1.forward(inputs)
    end = time.time()
    print('Time taken for forward prop on 1 stream: (sequentially)',end-start)
    
    start = time.time()
    outputs = net1.forward(inputs)
    outputs = net2.forward(inputs)
    end = time.time()
    print('Time taken for forward prop on 2 stream: (sequentially)',end-start)
    
    p1 = mp.Process(target=net1.forward, args=(inputs,))
    p2 = mp.Process(target=net2.forward, args=(inputs,))
    start = time.time()
    p1.start()
    p2.start()
    p1.join()
    p2.join()
    end = time.time()
    print('Time taken for forward prop on 2 stream: (parallel)',end-start)

这是输出：

Time taken for forward prop on 1 stream: (sequentially) 0.08776640892028809
Time taken for forward prop on 2 stream: (sequentially) 0.15159368515014648
Time taken for forward prop on 2 stream: (parallel) 3.8684606552124023

可以看出，前向道具是按顺序执行的，请问如何使两个网络的前向传播同时执行？

Answer 1

我的目标是使推理过程尽可能高效，因此我希望使 2 个不同的流同时运行。

您的代码以

开头

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

所以我假设您的目标是使用 CUDA，而您从未将数据或模型移动到 device 这一事实是疏忽。

torch.multiprocessing 创建多个进程，multiple processes have different default CUDA streams。

但是，运行多个操作系统线程或进程不会导致加速，因为 CUDA 在任何给定设备上实际上一次只运行其中一个 (1)

此外，根据我的经验，即使您在同一进程中使用了多个 CUDA 流，这也不会加速计算。这是一个通用的 CUDA 限制 (1)，与 PyTorch 无关。不过，流对于重叠 IO 和计算很有用。例如，您可以有一个 CUDA 流将数据复制到 GPU，而另一个正在对前一批进行推理。

---

(1) 这似乎随着 Ampere 架构而改变，因为它支持“多实例 GPU”。