在PyTorch中进行模型迁移学习通常需要以下步骤:
import torch
import torchvision.models as models
pretrained_model = models.resnet18(pretrained=True)
pretrained_model.fc = nn.Linear(pretrained_model.fc.in_features, num_classes)
for param in pretrained_model.parameters():
param.requires_grad = False
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(pretrained_model.fc.parameters(), lr=0.001)
for epoch in range(num_epochs):
for images, labels in dataloader:
optimizer.zero_grad()
outputs = pretrained_model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
通过以上步骤,你可以在PyTorch中进行模型迁移学习。你可以根据具体的任务需求对以上步骤进行调整和扩展。