如果在 Colab 中更改 FC 层，Pytorch Resnet 模型错误

Question

如果我只是从 Colab 中的 Pytorch 导入 Resnet 模型，并使用它来训练我的数据集，则没有问题。但是，当我尝试更改最后一个 FC 层以将输出特征从 1000 更改为 9，即我的数据集的类个数时，出现以下错误。

RuntimeError: Tensor for 'out' is on CPU, Tensor for argument #1 'self' is on CPU, but expected them to be on GPU (while checking arguments for addmm)

工作版本：

import torchvision.models as models
#model = Net()
model=models.resnet18(pretrained=True)

# defining the optimizer
optimizer = Adam(model.parameters(), lr=0.07)
# defining the loss function
criterion = CrossEntropyLoss()
# checking if GPU is available
if torch.cuda.is_available():
    model = model.cuda()
    criterion = criterion.cuda()

有错误的版本：

import torchvision.models as models
#model = Net()
model=models.resnet18(pretrained=True)

# defining the optimizer
optimizer = Adam(model.parameters(), lr=0.07)
# defining the loss function
criterion = CrossEntropyLoss()
# checking if GPU is available
if torch.cuda.is_available():
    model = model.cuda()
    criterion = criterion.cuda()
model.fc = torch.nn.Linear(512, 9)

错误发生在训练发生的阶段，又名

outputs = model(images)

我应该如何解决这个问题？

Answer 1

简单的错误，在将模型声明为cuda之前应该实例化fc层。即

model=models.resnet18(pretrained=True)
model.fc = torch.nn.Linear(512, 9)
if torch.cuda.is_available():
    model = model.cuda()

如果在 Colab 中更改 FC 层，Pytorch Resnet 模型错误

Pytorch Resnet model error if FC layer is changed in Colab

resnet

pytorch

google-colaboratory