时间戳

首页 > 解决方案 > 如果在 Colab 中更改 FC 层，则 Pytorch Resnet 模型错误

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

问题描述

如果我只是从 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)

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

标签： pytorchgoogle-colaboratoryresnet

解决方案

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

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

推荐阅读