C++libtorch 多线程多 GPU 推理问题

[justou]

遇到一个多 GPU 推理问题，模型是 TorchScript ，我想利用两块显卡并行推理，想法很直接：

1. 生成两个独立模型，model0 发送到 GPU0 上，model1 发送到 GPU1 上
2. 将一个 batch 的数据均分成两份，sub_batch0 发送到 GPU0 上, sub_batch1 发送到 GPU1 上
3. 在两个 c++线程中（std::thread）进行推理: model0(sub_batch0), model1(sub_batch1)并行推理后收集最终结果。

发现这样做有严重的性能问题，假如输入 batch=24,

1. 在单 GPU 上，显存占用~6G ，核心利用率~100%，推理速度~33fps
2. 在双 GPU 上（每个 GPU 获得大小为 12 的 sub_batch ），结果每张卡显存占用~15G （总 16GB ）, 核心占用率不到 50%，推理速度~44fps

按理说，在双 GPU 上，可以让 batch=48 ，从而每个显卡获的 sub_batch=24, 每个卡显存占用~6G, 核心利用率~100%，总的推理速度~66fps. 实际上 batch=48 直接爆显存了。

可能以上很直接的想法不符合 libtorch 或者说底层 cuda 的工作原理。请问下有朋友遇到过相同问题么？请教下如何解决。（也许需要用torch::nn::parallel::data_parallel来利用多 GPU ，我稍后尝试下）

[justou]

经过一些测试和查找资料，发现单个进程内，cuda 似乎无法充分利用多显卡，可能是单个进程 cuda 能利用的数据带宽有限，得一个进程一个模型对应一个显卡。

例如有 4 个显卡，主进程启动 4 个子进程，每个子进程一个推理模型，主进程源源不断地将数据一个一个 batch 的送到进程共享队列，每个推理子进程一边从队列获取数据缓存一边推理一边将结果发送回主线程汇总，这些都异步进行，主要目的是让各个显卡尽量不要空闲，以此将利用率最大化。

最开始尝试将大量输入数据均摊到各个子进程，然后等他们完成，再获取结果，发现这种方式显卡利用率不高，主要是有些进程很快完成推理，有些进程则很慢，这种方式永远是在等最慢的那个进程，而其它进程早就空闲下来了。这种方式得到的 fps 跟单个进程多个线程的差不多，而通过上面这种流水线处理方式，fps 按显卡个数倍增。

[allegory]

记得发后续

[ypw]

给每个推理线程创建独立的 stream 试试。

参考链接： https://pytorch.org/cppdocs/notes/tensor_cuda_stream.html

[justou]

@ypw 按理说应该是可行的, 只要每个线程都能在对应的设备上分配独立的 cuda context, 彼此互不影响. 后来我的应用要同时兼容 TorchScript 和 PyTorch model, 所以只能用进程了