当模型在设备之间移动时，向后（输入=）不起作用

Question

import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.weight_mul = nn.Parameter(torch.randn(D,))
        self.weight = nn.Parameter(torch.randn(D,))

    def forward(self, x):
        x = x * self.temp_weight
        return x

D = 5

x = torch.randn(D,).cuda()
model = Model()
model.cuda()
model.temp_weight = model.weight * model.weight_mul
model.cpu(); model.cuda()
output = model(x)
output.sum().backward(inputs=[model.weight, model.weight_mul])

print(model.weight.grad)
print(model.weight_mul.grad)

令我惊讶的是，

.grad

没有。有两种方法可以让

backward()

发挥作用。一是去掉

inputs

到

backward()

，二是去掉

model.cpu(); model.cuda()

。但为什么呢？

Answer 1

我认为发生这种情况是因为你在没有更新autograd引擎的情况下移动张量。

考虑这个例子：

D = 5
x = torch.ones((D,), device="cuda")
a = 2 * torch.ones((D,), device="cuda", requires_grad=True)
q = a # keep a reference to a
b = 3 * torch.ones((D,), device="cuda", requires_grad=True)
y = a*x + b
s = y.sum()


a = a.cpu()
a = a.cuda()

s.backward(inputs=[a])
print(a.grad) # --> None
s.backward(inputs=[q])
print(a.grad) # --> None
print(q.grad) # --> tensor([1., 1., 1., 1., 1.], device='cuda:0')

因此，当您写入

a = a.cpu()

或

a = a.cuda()

时，当发生移动到不同的设备内存时，您会获得

a

的副本。但是，反向传播的基础图不会自动更新以反映新内存类型中张量的新地址。您可以看到这是正确的，因为计算相对于

q

的梯度（指的是

a

的原始位置）按预期工作。

在您的情况下，当您调用

model.cpu()

或

model.cuda()

时，作为

nn.Parameter

属性的所有

model

对象都会发生同样的情况。由于

self.temp_weight

未更新以反映

model.weight

的新位置，因此 autograd 的行为会异常。

但是，请注意，这只是完整图片的一部分，因为如果您尝试在

nn.Parameter

中保存对

model

的引用，例如

q = model.weight

这不起作用。很明显，除了我上面概述的逻辑之外，还发生了其他事情。