遇到一个多 GPU 推理问题,模型是 TorchScript ,我想利用两块显卡并行推理,想法很直接:
- 生成两个独立模型,model0 发送到 GPU0 上,model1 发送到 GPU1 上
- 将一个 batch 的数据均分成两份,sub_batch0 发送到 GPU0 上, sub_batch1 发送到 GPU1 上
- 在两个 c++线程中(
std::thread)进行推理:model0(sub_batch0),model1(sub_batch1)并行推理后收集最终结果。
发现这样做有严重的性能问题,假如输入 batch=24,
- 在单 GPU 上,显存占用~6G ,核心利用率~100%,推理速度~33fps
- 在双 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 ,我稍后尝试下)
