在Torch中進行圖像分類任務通常涉及以下步驟:
準備數據集:首先,需要準備包含訓練和測試圖像的數據集。可以使用Torch的image
庫來加載和處理圖像數據。
定義模型架構:接下來,需要定義一個適合圖像分類任務的模型架構。可以使用Torch提供的預訓練模型,如VGG、ResNet、DenseNet等,也可以自定義模型架構。
定義損失函數:為了訓練模型,需要定義一個損失函數來衡量模型預測與真實標簽之間的差異。常用的損失函數包括交叉熵損失函數。
訓練模型:使用訓練集對模型進行訓練。可以使用Torch提供的nn
模塊來構建模型,并使用optim
模塊來定義優化器進行參數更新。
評估模型性能:使用測試集對訓練好的模型進行評估,計算模型在測試集上的準確率等性能指標。
下面是一個簡單的示例代碼,演示如何在Torch中進行圖像分類任務:
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import models, datasets, transforms
# 準備數據集
transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
train_dataset = datasets.ImageFolder('path/to/train/dataset', transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True)
test_dataset = datasets.ImageFolder('path/to/test/dataset', transform=transform)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=32, shuffle=False)
# 定義模型架構
model = models.resnet18(pretrained=True)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, len(train_dataset.classes))
# 定義損失函數和優化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
# 訓練模型
model.train()
for epoch in range(10):
for inputs, labels in train_loader:
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# 評估模型性能
model.eval()
correct = 0
total = 0
with torch.no_grad():
for inputs, labels in test_loader:
outputs = model(inputs)
_, predicted = torch.max(outputs, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print('Test accuracy: {} %'.format(100 * correct / total))
在這個示例中,我們使用了預訓練的ResNet-18模型進行圖像分類任務,使用ImageNet數據集進行預訓練。我們定義了一個簡單的訓練循環來訓練模型,并在測試集上評估模型性能。最后,我們輸出了模型在測試集上的準確率。