在PyTorch中進行遷移學習通常包括以下步驟:
torchvision.models
中的模型來加載預訓練模型。import torchvision.models as models
model = models.resnet18(pretrained=True)
model.fc = nn.Linear(model.fc.in_features, num_classes)
for param in model.parameters():
param.requires_grad = False
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001)
for epoch in range(num_epochs):
for images, labels in dataloader:
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
以上是在PyTorch中進行遷移學習的基本步驟,根據具體的任務和數據集可以對模型進行更多的調整和優化。