使用PyTorch DataParallel多GPU训练时程序崩溃的问题求助

问题拆解 #

• 参数显示翻倍：这不是实际参数翻倍，是torchsummary的统计特性导致的——DataParallel会把模型复制到每个GPU，torchsummary会将所有GPU上的模型参数累加统计，实际训练时参数是共享的，属于正常显示现象。
• GPU占满崩溃：本质是内存溢出（OOM）引发的进程崩溃，多GPU场景下的常见诱因包括内存拷贝顺序错误、Batch Size设置过大等。

具体解决步骤 #

1. 修正模型移动到GPU的顺序 #

先将模型移到目标GPU，再封装DataParallel，减少不必要的内存拷贝开销：

Inceptionv4_ = Inceptionv4()
# 先把模型移到目标设备
Inceptionv4_ = Inceptionv4_.to(device)
if tc.cuda.device_count() > 1:
    print(tc.cuda.device_count())
    Inceptionv4_ = nn.DataParallel(Inceptionv4_)

optimizer = tc.optim.SGD(Inceptionv4_.parameters(), lr=0.005)
# 查看真实参数需调用模型的module属性
if tc.cuda.device_count() > 1:
    print(summary(Inceptionv4_.module, (3, 224, 224)))
else:
    print(summary(Inceptionv4_, (3, 224, 224)))

2. 合理调整Batch Size #

多GPU训练时，每个GPU会分配到总Batch Size / GPU数量的数据，不要直接将单GPU的Batch Size乘以GPU数，建议调整为单GPU训练时的1.5~2倍，避免内存过载。

3. 优化数据加载配置 #

• 给DataLoader添加pin_memory=True，减少CPU到GPU的数据拷贝开销；
• 暂时关闭复杂的实时数据增强操作，降低内存占用。

4. 清理GPU内存碎片 #

训练前手动清理GPU缓存，释放碎片内存：

tc.cuda.empty_cache()

额外优化方案 #

如果主GPU（默认cuda:0）内存紧张，可以通过device_ids参数指定其他GPU作为主设备：

# 指定cuda:1为主GPU，cuda:0为辅助GPU
Inceptionv4_ = nn.DataParallel(Inceptionv4_, device_ids=[1, 0])

内容的提问来源于stack exchange，提问作者Uwe.Schneider