在PyTorch中實現遷移學習通常包括以下步驟:
import torch
import torchvision.models as models
model = models.resnet18(pretrained=True)
n_features = model.fc.in_features
model.fc = torch.nn.Linear(n_features, num_classes) # num_classes為新任務的類別數
for param in model.parameters():
param.requires_grad = False
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
for epoch in range(num_epochs):
for inputs, labels in dataloader:
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
這樣就完成了遷移學習的實現過程。通過以上步驟,你可以利用預訓練的模型在新任務上快速進行模型訓練。