溫馨提示×

如何在PyTorch中使用預(yù)訓(xùn)練的模型進行遷移學(xué)習(xí)

小樊
107
2024-03-05 19:33:06
欄目: 編程語言

在PyTorch中使用預(yù)訓(xùn)練的模型進行遷移學(xué)習(xí)可以通過以下步驟實現(xiàn)：

  1. 加載預(yù)訓(xùn)練模型：首先，使用torchvision.models模塊中提供的預(yù)訓(xùn)練模型來加載已經(jīng)訓(xùn)練好的模型，例如ResNet、VGG等。
import torchvision.models as models

# Load pre-trained ResNet-50 model
model = models.resnet50(pretrained=True)
  1. 修改模型的最后一層：由于遷移學(xué)習(xí)通常涉及到不同的任務(wù)，需要將預(yù)訓(xùn)練模型的最后一層替換為新的全連接層，以適應(yīng)新的任務(wù)要求。
import torch.nn as nn

# Modify the last layer of the model
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, num_classes) # num_classes為新任務(wù)的類別數(shù)
  1. 設(shè)置優(yōu)化器和損失函數(shù)：根據(jù)新任務(wù)的要求設(shè)置優(yōu)化器和損失函數(shù)。
import torch.optim as optim

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
  1. 訓(xùn)練模型：使用新的數(shù)據(jù)集對模型進行訓(xùn)練。
# 訓(xùn)練代碼

通過以上步驟，您可以在PyTorch中使用預(yù)訓(xùn)練的模型進行遷移學(xué)習(xí)。

0 踩