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

242 天前
 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 ,我稍后尝试下)

872 次点击
所在节点    C++
2 条回复
justou
232 天前
经过一些测试和查找资料,发现单个进程内,cuda 似乎无法充分利用多显卡,可能是单个进程 cuda 能利用的数据带宽有限,得一个进程一个模型对应一个显卡。

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

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

这是一个专为移动设备优化的页面(即为了让你能够在 Google 搜索结果里秒开这个页面),如果你希望参与 V2EX 社区的讨论,你可以继续到 V2EX 上打开本讨论主题的完整版本。

https://www.v2ex.com/t/972352

V2EX 是创意工作者们的社区,是一个分享自己正在做的有趣事物、交流想法,可以遇见新朋友甚至新机会的地方。

V2EX is a community of developers, designers and creative people.

© 2021 V2EX