PyTorch多GPU推理采用数据并行时第二块GPU始终空闲问题咨询

• 输入Batch size不符合拆分要求
  nn.DataParallel默认对输入张量的第0维度按GPU数量均分拆分，若单轮推理的batch size小于GPU数量，拆分后第二块卡无计算任务，会处于空闲状态。请将推理batch size调整为不小于2的数值，优先设置为GPU数量的整数倍以获得最优负载。
• 模型封装与设备迁移顺序错误
  当前代码执行顺序不符合DataParallel的调用要求，封装DataParallel前需要先将模型迁移到主GPU（device_ids[0]对应的GPU），封装后无需再额外调用to(device)，修正后的推理初始化代码如下：

import torch
import torch.nn as nn

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
n_gpu = torch.cuda.device_count()

if torch.cuda.is_available():
    # 先将基础模型迁移到主GPU
    net = net.to(device)
    # 加载模型权重
    net.load_state_dict(torch.load(model_dir, map_location=device))
    # 多卡场景下封装DataParallel
    if n_gpu > 1:
        net = nn.DataParallel(net, device_ids = [0, 1])

• 分布式训练保存的权重加载逻辑错误
  如果你使用的模型权重是训练阶段用DataParallel封装后保存的，直接加载到未封装的普通模型会出现权重key前缀不匹配的问题，导致模型加载异常，进而影响多卡调度。可以通过以下方式处理权重后加载：

state_dict = torch.load(model_dir, map_location=device)
# 去掉权重key的module前缀
fixed_state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
net.load_state_dict(fixed_state_dict)

• 推理时未关闭梯度计算
  推理时如果不关闭梯度计算，会额外占用主卡显存，也可能干扰多卡任务调度。推理执行阶段请套入torch.no_grad()上下文管理器，示例如下：

with torch.no_grad():
    output = net(input_data)

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