使用 PyTorch 训练 VGG16 模型进行图像分类

Question

我正在使用 PyTorch 进行图像分类。

我编写了以下使用简单线性模型的训练函数：

criterion = nn.CrossEntropyLoss()
def train(model, dataloader, epoch):
model.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
running_loss, running_acc = 0., 0.
loss_history = []
accuracy_history = [](data_train):.2f}%")

for i in range(1, epoch + 1):
  model.train()
  for inputs, targets in dataloader:
      inputs, targets = inputs.to(device), targets.to(device)
      outputs = model(inputs)
      loss = criterion(outputs, targets)
      optimizer.zero_grad()
      loss.backward()
      optimizer.step()
      preds = torch.argmax(outputs, 1)
      running_loss += loss.item()
      running_acc += torch.sum(preds == targets).item()

  print(f"[TRAIN epoch {i}] Loss: {running_loss/len(data_train):.2f} Acc: {100 * running_acc/len
 

我有预训练的 VGG16 模型，我想更改其最后一层的权重：

model_vgg = models.vgg16(weights='DEFAULT')
model_vgg.classifier[6] = nn.Linear(4096, 2)

for param in model_vgg.parameters():
    param.requires_grad = False
model_vgg.classifier[-1].requires_grad = True

train(model_vgg, train_loader, 2)

但是，在训练时出现以下错误：

RuntimeError                              Traceback (most recent call last)

<timed eval> in <module>

<ipython-input-27-1f64686a5cfd> in train(model, dataloader, epoch)
     39           loss = criterion(outputs, targets)
     40           optimizer.zero_grad()
---> 41           loss.backward()
     42           optimizer.step()
     43           preds = torch.argmax(outputs, 1)

/usr/local/lib/python3.10/dist-packages/torch/autograd/__init__.py in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)
--> 251     Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
    252         tensors,
    253         grad_tensors_,

RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

我该如何解决这个问题？

Answer 1

您需要确保仅对最后一层的参数启用渐变。将

model_vgg.classifier[-1].requires_grad = True

替换为以下代码片段

for param in model_vgg.classifier[-1].parameters():
    param.requires_grad = True