目录
一、项目背景
随着图像数据的海量增长,对图像进行快速准确的分类成为了一项重要的任务。在本项目中,我们选择了常见的图像分类问题,例如区分猫和狗的图像,旨在通过深度学习技术,特别是基于卷积神经网络的方法,实现对这类图像的自动分类。
考虑到实际应用中的需求,我们需要一个能够在不同设备上运行,并且具有较高准确性和泛化能力的模型。通过使用流行的深度学习框架 PyTorch ,我们能够利用其丰富的功能和优化算法,来构建和训练我们的模型。
我们收集了大量的猫和狗的图像数据集,并对其进行了预处理和划分,分为训练集、验证集和测试集。通过精心设计的图像变换和数据增强策略,提高模型的鲁棒性和泛化能力
二、环境与库准备
首先,我们导入了一系列必要的库,包括 torch
相关的模块用于模型构建和训练,torchvision
用于数据处理和模型加载,以及一些数据处理和可视化的库。这些库为我们提供了构建、训练和评估模型所需的各种功能和工具。
import torch.optim as optim
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.optim
import torch.utils.data
import torch.utils.data.distributed
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torchvision.models
from torch.autograd import Variable
from tqdm import tqdm
import os
from PIL import Image
import cv2
三、设备选择
接下来,通过以下代码选择计算设备:
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
DEVICE
这行代码会自动检测是否有可用的 GPU,如果有则使用 GPU('cuda')进行计算,否则使用 CPU。
四、数据预处理
定义了 transform111
字典,其中包含了针对训练数据和验证数据的不同预处理操作,如随机裁剪、旋转、翻转、标准化等,以增加数据的多样性和模型的泛化能力。
transform111 = {
'train': transforms.Compose([
transforms.RandomResizedCrop(size=256, scale=(0.8, 1.0)),
transforms.RandomRotation(degrees=15),
transforms.RandomHorizontalFlip(),
transforms.CenterCrop(size=224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])
]),
'val': transforms.Compose([
transforms.Resize(size=256),
transforms.CenterCrop(size=224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])
])
}
五、数据读取与准备
读取训练和验证数据的文件夹,并使用 ImageFolder
类进行组织。然后,计算训练和验证数据的大小,创建数据加载器,设置批处理大小和工作线程数,为后续的模型训练和验证做好准备。
import shutil
modellr = 1e-3
BATCH_SIZE = 64
EPOCHS = 4
# 读取数据
dataset_train = datasets.ImageFolder('dataset/train', transform111)
print(dataset_train.imgs)
# 对应文件夹的label
print(dataset_train.class_to_idx)
dataset_test = datasets.ImageFolder('dataset/val', transform111)
# 对应文件夹的label
print(dataset_test.class_to_idx)
dataset = './dataset'
train_directory = os.path.join(dataset, 'train')
valid_directory = os.path.join(dataset, 'val')
batch_size = 32
num_classes = 6
print(train_directory)
data = {
'train': datasets.ImageFolder(root=train_directory, transform=transform111['train']),
'val': datasets.ImageFolder(root=valid_directory, transform=transform111['val'])
}
train_data_size = len(data['train'])
valid_data_size = len(data['val'])
train_loader = torch.utils.data.DataLoader(data['train'], batch_size=batch_size, shuffle=True, num_workers=8)
test_loader = torch.utils.data.DataLoader(data['val'], batch_size=batch_size, shuffle=True, num_workers=8)
print(train_data_size, valid_data_size)
六、模型构建与优化
- 模型定义
使用 ResNet50
模型结构,但不加载预训练权重。
# 下载预训练模型
model = torchvision.models.resnet50(pretrained=False) # 不使用训练好的预训练模型
model
2.全连接层调整
criterion = nn.CrossEntropyLoss()
num_ftrs = model.fc.in_features
model.fc = nn.Sequential(
nn.Linear(num_ftrs, 256),
nn.ReLU(),
nn.Dropout(0.4),
nn.Linear(256, 6), # 猫狗是二分类,故是 2
nn.LogSoftmax(dim=1)
)
model.to(DEVICE)
3.优化器选择与学习率设置
选择 Adam
优化器,并定义了学习率调整函数,根据训练轮数动态降低学习率。
optimizer = optim.Adam(model.parameters(), lr=modellr)
def adjust_learning_rate(optimizer, epoch):
modellrnew = modellr * (0.1 ** (epoch // 50))
print("lr:", modellrnew)
for param_group in optimizer.param_groups:
param_group['lr'] = modellrnew
七、训练与验证过程
-
定义了
train
函数进行模型的训练,在每个批次中计算损失、进行反向传播和参数更新,并定期打印训练信息和计算平均准确率与损失。def train(model, device, train_loader, optimizer, epoch): model.train() sum_loss = 0 train_acc = 0 total_num = len(train_loader.dataset) for batch_idx, (data, target) in enumerate(tqdm(train_loader)): data, target = Variable(data).to(device), Variable(target).to(device) im = Variable(data) output = model(data) loss = criterion(output, target) optimizer.zero_grad() loss.backward() optimizer.step() print_loss = loss.data.item() sum_loss += print_loss out_t = output.argmax(dim=1) # 取出预测的最大值 num_correct = (out_t == target).sum().item() acc = num_correct / im.shape[0] train_acc += acc if (batch_idx + 1) % 50 == 0: print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( epoch, (batch_idx + 1) * len(data), len(train_loader.dataset), 100. * (batch_idx + 1) / len(train_loader), loss.item())) ave_loss = sum_loss / len(train_loader) ave_acc = train_acc / len(train_loader) print('epoch:{}, train_acc: {}, loss:{}'.format(epoch, ave_acc, ave_loss)) return ave_acc, ave_loss
- 定义了
val
函数进行模型的验证,计算验证集上的损失和准确率。def val(model, device, test_loader): model.eval() test_loss = 0 correct = 0 total_num = len(test_loader.dataset) with torch.no_grad(): for data, target in test_loader: data, target = Variable(data).to(device), Variable(target).to(device) output = model(data) loss = criterion(output, target) _, pred = torch.max(output.data, 1) correct += torch.sum(pred == target) print_loss = loss.data.item() test_loss += print_loss correct = correct.data.item() acc = correct / total_num avgloss = test_loss / len(test_loader) print('Val set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format( avgloss, correct, len(test_loader.dataset), 100 * acc)) return acc, avgloss
在训练过程中,通过循环遍历每个训练轮次,调用学习率调整函数、训练函数和验证函数,记录训练和验证的准确率与损失。
然后通过以下代码进行训练和验证
train_acc_list, train_loss_list, val_acc_list, val_loss_list = [], [], [], []
for epoch in range(1, EPOCHS + 1):
adjust_learning_rate(optimizer, epoch)
train_acc, train_loss = train(model, DEVICE, train_loader, optimizer, epoch)
val_acc, val_loss = val(model, DEVICE, test_loader)
train_acc_list.append(train_acc)
val_acc_list.append(val_acc)
train_loss_list.append(train_loss)
val_loss_list.append(val_loss)
八、模型保存与可视化
训练完成后,保存训练好的模型,并使用 matplotlib
绘制训练和验证的准确率与损失曲线,以便直观地观察模型的性能。
torch.save(model,'model.pth')
import matplotlib.pyplot as plt
epochs_range = range(EPOCHS)
print(epochs_range, train_acc_list)
plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)
plt.plot(epochs_range, train_acc_list, label='Training Accuracy')
plt.plot(epochs_range, val_acc_list, label='Validation Accuracy')
plt.legend(loc='lower right')
plt.title('Training and Validation Accuracy')
plt.subplot(1, 2, 2)
plt.plot(epochs_range, train_loss_list, label='Training Loss')
plt.plot(epochs_range, val_loss_list, label='Validation Loss')
plt.legend(loc='upper right')
plt.title('Training and Validation Loss')
plt.savefig('./acc-loss.jpg')
预测准确率随着预测次数而提高
九、模型测试
最后,定义了测试数据的预处理操作,加载保存的模型,对测试集进行预测,并输出每个测试样本的图像名称、真实标签和预测标签。
import torch.utils.data.distributed
import torchvision.transforms as transforms
import torchvision.datasets as datasets
from torch.autograd import Variable
from pathlib import Path
classes = ('cat', 'dog')
transform_test = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = torch.load("model.pth")
model.eval()
model.to(DEVICE)
dataset_test = datasets.ImageFolder('dataset/test', transform_test)
print(len(dataset_test))
# 对应文件夹的label
y_true, y_sore = [], []
for index in range(len(dataset_test)):
item = dataset_test[index]
img, label = item
img.unsqueeze_(0)
data = Variable(img).to(DEVICE)
output = model(data)
_, pred = torch.max(output.data, 1)
y_true.append(label)
y_sore.append(pred.data.item())
print('Image Name:{}, label:{}, predict:{}'.format(dataset_test.imgs[index][0], classes[label], classes[pred.data.item()]))
index += 1
图像分类预测结果的记录