Pytorch 自定义模型自动存储在 cuda 中

Question

我像这样构建了一个自定义 NN 模型：

class MyNNet(torch.nn.Module):
  
  def __init__(self, inp_dim, n_classes):
    super(MyNNet, self).__init__()
    self.flat = torch.nn.Flatten()
    self.l1 = torch.nn.Linear(inp_dim * inp_dim, 32)
    self.l2 = torch.nn.Linear(32, 16)
    self.l3 = torch.nn.Linear(16, n_classes)
  
  def forward(self, X):
    out = self.flat(X)
    out = F.relu(self.l1(out))
    out = F.relu(self.l2(out))
    return self.l3(out)

还有一个更新模型参数的简单训练脚本：

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = MyNNet(28, 10)
model.to(device)

optimizer = torch.optim.Adam(model.parameters())
loss = torch.nn.CrossEntropyLoss()

epochs = 20

for e in range(epochs):

  train_l = 0.
  for i, (s, c) in enumerate(train_loader):
    
    s.to(device)
    c.to(device)
    y_hat = model(s)

    l = loss(y_hat, c)
    train_l += l

    l.backward()
    optimizer.step()
    optimizer.zero_grad()

  print(f'Epoch: {e}, AvgLoss: {train_l / len(train_loader)}')

在脚本中，我将模型存储到 cuda，因此我对每批数据集 (MNIST) 进行处理。但是出现以下错误：Expected all tensors to be on the same device, but found at least two devices 但是当我评论model.to(device) 时，脚本就起作用了。这是否意味着 PyTorch 会自动将自定义模型存储到 cuda 中？
谢谢。

Answer 1

与Modules（.to(...) 就地工作）不同，将Tensors 移动到设备时，您需要重新分配它们：

s = s.to(device)
c = c.to(device)