pytorch使用DataParallel时遇到的几个问题

原创已于 2022-03-23 17:57:29 修改 · 3.1k 阅读

3 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #深度学习 #python

于 2021-11-15 15:58:05 首次发布

pytorch 专栏收录该内容

47 篇文章

订阅专栏

这篇博客讨论了在使用DataParallel进行多GPU训练时遇到的问题。当使用两块GPU，总损失会自动取平均值，但自定义损失可能不会，导致`item()`调用错误。解决方法是在自定义损失上应用`torch.mean()`来获取平均值，然后再使用`item()`转换为浮点数。此外，文章也强调了DataParallel如何将批量数据分割到多个GPU上进行并行处理。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

import Model
import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = Model(input)
model.to(device)
model = torch.nn.DataParallel(model)
...

1、问题1：
ValueError: only one element tensors can be converted to Python scalars
原因：
（1）
单块卡返回loss是一个tensor，使用loss.item()得到float类型的值；
（2）
DataParallel使用了两块gpu，此时总的loss会自动计算平均值，
但是其他自定义的loss不会自动计算平均值，因此返回的是一个list，如果对list使用item（）函数，是会报错的。
改进：加一个torch.mean（），手动计算其他loss的平均值，然后使用item()取出浮点类型的数值。
值得注意的是，使用dataparalle是数据并行的，也就是说，输入为[batch_size, seq_len]的数据，当加载到两块gpu上是，每块gpu会分别处理[batch_size//2, seq_len]的数据。