使用 CUDA 张量和 CPU 张量的 PyTorch 前向传递，同时保留梯度

Question

我有一个由神经模型生成的向量

a

，需要与一个巨大的矩阵

M

交互。由于

M

很大，我必须在cpu设备中进行计算。这样的话，我想知道在cuda设备上是否可以保留梯度并向后倾斜。

示例如下：

a_cuda = torch.randn([1, 512], requires_grad=True).to("cuda")
a_cpu = torch.randn([1, 512], requires_grad=True).to("cpu")

M = torch.randn([512, 100000], requires_grad=False) # loaded in cpu device, dont need update

out_cuda = (a_cuda.cpu() @ M).sum()
out_cuda.backward()

out_cpu = (a_cpu @ M).sum()
out_cpu.backward()

print(a_cuda.grad) # None
print(a_cpu.grad)

我正在寻找解决方案，使得

a_cuda.grad

具有与

a_cpu.grad

相同的梯度。

Answer 1

当然。

试试这个：

a_cuda = torch.randn([1, 512], requires_grad=True, device='cuda')
a_cpu = torch.randn([1, 512], requires_grad=True)

M = torch.randn([512, 100000], requires_grad=False) # loaded in cpu device, dont need update

out_cuda = (a_cuda.cpu() @ M).sum()
out_cuda.backward()

out_cpu = (a_cpu @ M).sum()
out_cpu.backward()

print(a_cuda.grad)
print(a_cpu.grad)