遷移學習(Transfer learning)是一種機器學習技術,通過將在一個任務上學到的知識遷移到另一個相關任務中,可以加快模型訓練過程并提高模型性能。在本文中,我們將介紹基于Torch框架的遷移學習應用案例。
案例背景:
假設我們有一個已經訓練好的圖像分類模型,可以識別各種動物的圖像。現在我們想要基于這個模型,構建一個新的模型,用于識別植物的圖像。由于動物和植物的圖像具有一定的相似性,我們可以利用已有的動物分類模型進行遷移學習,加速訓練并提高模型準確率。
實現步驟:
1. 加載預訓練的動物分類模型(如ResNet、VGG等)并凍結模型參數。
2. 替換模型的最后一層全連接層,將輸出類別數調整為植物類別數。
3. 定義新的數據集,并加載植物圖像數據集。
4. 設置損失函數和優化器。
5. 進行模型訓練,僅更新新添加的全連接層參數。
6. 對模型進行微調(可選),解凍部分或全部模型參數,繼續訓練。
代碼示例(基于Torch):
```python
import torch
import torchvision
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms, models
# 加載預訓練的動物分類模型
model = models.resnet18(pretrained=True)
# 凍結模型參數
for param in model.parameters():
param.requires_grad = False
# 替換最后一層全連接層
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, num_classes) # num_classes為植物類別數
# 定義損失函數和優化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.fc.parameters(), lr=0.001)
# 加載植物圖像數據集
data_transforms = {
'train': transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]),
'val': transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]),
}
image_datasets = {x: datasets.ImageFolder(f'path/to/{x}', data_transforms[x]) for x in ['train', 'val']}
dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=4, shuffle=True, num_workers=4) for x in ['train', 'val']}
# 訓練模型
def train_model(model, criterion, optimizer, num_epochs=5):
for epoch in range(num_epochs):
for inputs, labels in dataloaders['train']:
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# 進行模型微調
for param in model.parameters():
param.requires_grad = True
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
train_model(model, criterion, optimizer, num_epochs=5)
```
通過上述步驟和代碼示例,我們可以基于預訓練的動物分類模型,快速構建一個新的用于植物圖像分類的模型,并進行訓練和微調,以提高模型性能。這是一個簡單而有效的遷移學習應用案例,可以在實際項目中得到廣泛應用。