Pytorch | Pytorch框架中模型和数据的gpu和cpu模式:model.to(device), model.cuda(), model.cpu(), DataParallel

2022-09-21 13:28:50

背景介绍

我们在使用Pytorch训练时,模型和数据有可能加载在不同的设备上(gpu和cpu),在算梯度或者loss的时候,报错信息类似如下:

RuntimeError: Function AddBackward0 returned an invalid gradient at index 1 - expected type torch.cuda.FloatTensor but got torch.FloatTensor

RuntimeError: Expected object of device type cuda but got device type cpu for argument #2 'target' in call to _thnn_l1_loss_forward

这种报错通常有几种情况:

1. 数据在cpu上,模型在gpu上;

2. 数据在gpu上,模型在cpu上;

3. 指定层不支持gpu或cpu版本,但数据和模型都在gpu或cpu上。

解决方案

第1步:将模型和数据都移到同一设备上。

假设模型和数据定义如下:

model = Model()
input = dataloader()
output = model(input)

移动到gpu上:

# solution: 0
device = 'gpu'
model = model.to(device)
data = data.to(device)

# solution: 1
model = model.cuda()
data = data.cuda()

移动到cpu上:

# solution: 0
device = 'cpu'
model = model.to(device)
data = data.to(device)

# solution: 1
model = model.cpu()
data = data.cpu()

第二步:打印模型model和数据data在gpu上还是cpu上。

1. 通过判断模型model的参数是否在cuda上来判定模型是否在gpu上。

print('Is model on gpu: ', next(model.parameters()).is_cuda)

输出若是True,则model在gpu上;若是False,则model在cpu上。

2. 输出数据data的device字段。

print('data device: ', data.device)

输出gpu则在gpu上,输出cpu则在cpu上。

第三步:排查策略。

一般故事就到这里结束了。但是实际很多情况中,因为模型定义、数据加载复杂多样,很多时候还是会出现开头的报错。这个时候建议大家:

1. 在模型的定义开头处,设置统一的device变量。

有很多论坛里的解法建议大家添加了层之后,就逐层进行第一步里的设置。但是这样的操作很繁琐,而且可能会有很多冗余和缺漏。在模型定义之后,统一移到指定device上,然后通过第二步的输出,确保模型定义model和数据data都在指定的device上即可。

2. 逐步进行第二步的输出。

我遇到的情况是,我的模型定义后,参数已经在cpu上,数据也在cpu上,但是在训练过程中仍然报错。所以我在每个关键节点处都设置了第二步的打印策略,检查是不是中间处理步骤有设置出错。

3. 定位问题,修改代码。

DataParallel:自动拷贝到GPU上

前面提到的排查策略后,我发现在模型设置移到cpu之后(打印输出正确),仍然出现开头的报错。定位到问题就是在模型model定义后,还使用了DataParallel进行多卡(多gpu)的训练设置。

具体的官方定义在:data_parallel_tutorial

在使用了DataParallel api之后,模型model会被自动拷贝到可用gpu上。当需求是在cpu模式下进行训练,DataParallel就会自动把数据和模型都拷贝到gpu上,所以就有了开头的报错。

这个是在网上找到的解答:does-dataparallel-matters-in-cpu-mode

Ya, In CPU mode you cannot use DataParallel ().
Wrapping a module with DataParallel () simply copies the model over multiple GPUs and puts the results in device_ids[0] i.e. in 1st of the number of GPUs provided. If you are running in CPU mode, you should simply remove the DataParallel () wrapping. Regarding the documentation line “The code does not need to be changed in CPU-mode.” I guess it means the rest of code( other than model definition) need not be changed. But I am not sure about this and someone else can explain this part.

直白来说,在cpu模式下,我们不能使用DataParallel。该API会直接把模型拷贝到多块可用的gpu上,并把结果保存在第0块gpu上。若我们的模型运行在cpu模式下,我们需要移除DataParallel()。

在去掉了DataParallel之后,模型就可以和数据运行在cpu上了。

欢迎大家留言区指正~ 预祝大家炼丹顺利哈哈哈哈哈哈~

  • 作者:1LOVESJohnny
  • 原文链接:https://blog.csdn.net/iLOVEJohnny/article/details/106021547
    更新时间:2022-09-21 13:28:50