Pytorch：使用 CUDA 会阻止优化工作

Question

我有一个非常简单的优化：一条直线。代码如下：

use_gpu = torch.cuda.is_available()
learning_rate = 0.05
loss_function = nn.MSELoss()

train_inputs = torch.FloatTensor([1,2,3,4,5,6]).T.unsqueeze(0)
y_truth = torch.FloatTensor([10, 15, 20, 25, 30, 35]).unsqueeze(0)

W = torch.nn.Parameter(torch.rand(1), requires_grad=True)
b = torch.nn.Parameter(torch.rand(1), requires_grad=True)
optimizer = optim.Adam([b, W], lr=learning_rate)

# if use_gpu:
#     y_truth = y_truth.cuda()
#     W = W.cuda()
#     b = b.cuda()
#     train_inputs = train_inputs.cuda()



for epoch in range(1000):
    optimizer.zero_grad()
    y_preds = b + W * train_inputs
    loss = loss_function(y_truth, y_preds)
    loss.backward()
    optimizer.step()

    if epoch % 100 == 0:
        print(loss.data, W.data, b.data)

如果我不将数据放在 GPU 上，该代码可以正常工作。如果我取消注释 if use_gpu 块，代码会运行，但不会最小化任何内容并且变量不会更新。

我希望代码在 GPU 上也能以类似方式工作。知道发生了什么吗？

谢谢！

Answer 1

知道发生了什么吗？

是的，您正在训练的参数 W 和 b 保留在主机 (CPU) 上。

什么时候做的

W = W.cuda()
b = b.cuda()

您只是选择忽略正在优化的实际参数。

---

如果您希望为此使用 GPU，您可以尝试：

W = torch.nn.Parameter(torch.rand(1).cuda())
b = torch.nn.Parameter(torch.rand(1).cuda())

改为。