在PyTorch中使用預訓練的模型進行遷移學習可以通過以下步驟實現:
import torchvision.models as models
# Load pre-trained ResNet-50 model
model = models.resnet50(pretrained=True)
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為新任務的類別數
import torch.optim as optim
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
# 訓練代碼
通過以上步驟,您可以在PyTorch中使用預訓練的模型進行遷移學習。