PyTorch下Faster R-CNN配置DataParallel后仅单GPU工作问题排查

一、单GPU占用问题排查与修复 #

核心原因 #

你传入模型的第一个输入images是Python列表类型（每个元素为单张3维图片张量），nn.DataParallel默认仅对张量类型的第一个输入做第0维切分，非张量输入不会被切分，全部数据都会被发送到0号主卡，导致其余3张卡无数据、利用率为0。

修复步骤 #

• 1. 确认GPU可见性：在代码最开头添加以下逻辑，确保4张卡都被PyTorch识别

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2,3'
import torch
# 执行后确认输出为4，否则检查设备驱动和环境配置
print(torch.cuda.device_count())

• 2. 适配输入格式：将图片列表堆叠为4维batch张量后再传入模型，torchvision官方实现的Faster R-CNN原生支持该输入格式：

# 原有输入迁移逻辑改为堆叠操作
images = torch.stack([image.to(device) for image in images])
targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
loss_dict = model(images, targets)

• 3. 校验batch_size配置：传入DataLoader的总batch_size需要是单卡训练时batch_size的4倍，且满足你代码中batch_size % 4 == 0的校验规则。

可选优化方案 #

nn.DataParallel为单进程多线程实现，存在GIL瓶颈、主卡负载过高的问题，建议迁移到torch.nn.parallel.DistributedDataParallel（DDP）多进程训练框架，训练效率可提升30%以上，且对复杂输入格式的适配性更好。

二、标量拼接警告修复 #

你遇到的Was asked to gather along dimension 0, but all input tensors were scalars警告是因为每张卡返回的损失值为0维标量，DataParallel拼接多卡返回结果时找不到第0维触发，不影响训练正确性，消除警告可选择以下任意一种方法：

• 方法1：在模型返回损失字典前，对每个损失值升维

# 在Faster R-CNN模型返回loss_dict前添加如下逻辑
loss_dict = {k: v.unsqueeze(0) for k, v in loss_dict.items()}

• 方法2：合法过滤该警告

import warnings
warnings.filterwarnings("ignore", message="Was asked to gather along dimension 0, but all input tensors were scalars")

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