目录
1.数据集准备
收集大量的图像数据,包括各个类别的图像,确保每个类别都有足够的样本量。
1.1收集途径:
- 公共数据集(如ImageNet、CIFAR-10)
- 自行采集(通过网络爬虫、摄像头等)
- 商业数据集(购买或许可使用)
1.2 数据标注
-
确保每张图像都已标注,即图像所属的类别。
-
方式:通过手动标注或使用标注工具进行半自动标注。
1.3 数据划分
训练集 | 70-80% | 模型训练 |
验证集 | 10-15% | 调参和验证模型性能,避免过拟合 |
测试集 | 10-15% | 最终评估模型的泛化性能 |
2.数据预处理
2.1 图像大小调整
将所有图像调整为相同大小,通常是224x224像素,这是ResNet-50的输入大小要求。
from torchvision import transforms
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
])
2.2 数据增强(可选)
在训练集中进行数据增强,以增加数据的多样性,提高模型的泛化能力。
常见的方法:
- 随机裁剪
- 水平翻转
- 颜色抖动
- 旋转
data_augmentation = transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.5),
transforms.ToTensor(),
])
2.3 标准化
对图像数据进行标准化,使其具有相同的均值和标准差。
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
3.模型训练与验证
3.1 数据加载
使用PyTorch的DataLoader将数据加载到模型中。
from torchvision import datasets
from torch.utils.data import DataLoader
# 加载训练集
train_dataset = datasets.ImageFolder(root='path_to_train_data', transform=transform)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4)
# 加载验证集
val_dataset = datasets.ImageFolder(root='path_to_val_data', transform=transform)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4)
# 加载测试集
test_dataset = datasets.ImageFolder(root='path_to_test_data', transform=transform)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=4)
3.2 训练及验证
1.模型训练:在训练集上训练ResNet-50模型
-
损失函数:选择适合的损失函数,如交叉熵损失函数,用于分类任务。
-
优化器:选择优化算法,如Adam、SGD等,用于更新模型权重。
-
训练过程:将训练数据输入模型,进行前向传播计算输出,与实际标签计算损失,然后进行反向传播更新权重
2.验证:并使用验证集调整超参数和监控模型性能,避免过拟合。
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import models
# 加载预训练的ResNet-50模型
model = models.resnet50(pretrained=True)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, num_classes) # num_classes为分类类别数量
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
for epoch in range(num_epochs):
model.train()
running_loss = 0.0
for inputs, labels in train_loader:
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
# 验证
model.eval()
val_loss = 0.0
correct = 0
total = 0
with torch.no_grad():
for inputs, labels in val_loader:
outputs = model(inputs)
loss = criterion(outputs, labels)
val_loss += loss.item()
_, predicted = torch.max(outputs, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print(f'Epoch {epoch+1}/{num_epochs}, Loss: {running_loss/len(train_loader)}, Validation Loss: {val_loss/len(val_loader)}, Accuracy: {correct/total}')
4.模型测试
- 测试:在测试集上评估模型的最终性能
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(f'Test Accuracy: {correct/total}')