原因
可能在将需要求梯度的tensor放置在cuda上时出错,如
import torch
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
x = torch.rand(10, 3).to(device)
w = torch.rand(3, 1, requires_grad=True).to(device)
t = (x @ w).sum()
t.backward()
print(w.grad)
错误原因:对需要求梯度的张量做了一次to(device)导致叶节点发生变化
一种处理方式如下:
device = torch.device('cpu')
torch.manual_seed(123)
x = torch.rand(10, 3).to(device)
w = torch.rand(3, 1, requires_grad=True)
w1 = w.to(device)
t = (x @ w1).sum()
t.backward()
print(w.grad)
[4.1784],
[5.2859]])
这种方式看似work,实际上在写训练模块时,要将w1写在循环内,后续梯度求导就会出现问题,参数得不到更新
正确的处理方式
直接利用人家写好的接口就行啦,别花里胡哨的,.cuda() ,to(device)的。
import torch
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
torch.manual_seed(123)
x = torch.rand(10, 3).to(device)
w = torch.rand(3, 1, requires_grad=True, device=device)
t = (x @ w).sum()
t.backward()
print(w.grad)
[4.1784],
[5.2859]], device='cuda:0')
这样用就不会有什么问题了。
|