导入所需的Python库,包括图像处理、深度学习模型和数据加载
- import os
-
- import torch
-
- import torch.nn as nn
-
- import torch.optim as optim
-
- from torch.utils.data import Dataset, DataLoader
-
- from PIL import Image
-
- from torchvision import models, transforms
创建一个自定义的图像数据集类,用于加载和处理图像数据。
- class CustomImageDataset(Dataset):
-
- def __init__(self, main_dir, transform=None):
-
- self.main_dir = main_dir
-
- self.transform = transform
-
- self.files = []
-
- self.labels = []
-
- self.label_to_index = {}
-
- for index, label in enumerate(os.listdir(main_dir)):
-
- self.label_to_index[label] = index
-
- label_dir = os.path.join(main_dir, label)
-
- if os.path.isdir(label_dir):
-
- for file in os.listdir(label_dir):
-
- self.files.append(os.path.join(label_dir, file))
-
- self.labels.append(label)
-
-
-
- def __len__(self):
-
- return len(self.files)
-
-
-
- def __getitem__(self, idx):
-
- image = Image.open(self.files[idx])
-
- label = self.labels[idx]
-
- if self.transform:
-
- image = self.transform(image)
-
- return image, self.label_to_index[label]
定义一个数据转换过程,包括图像大小调整、随机翻转、旋转、转换为张量以及标准化
- transform = transforms.Compose([
-
- transforms.Resize((227, 227)), # AlexNet的输入图像大小
-
- transforms.RandomHorizontalFlip(), # 随机水平翻转
-
- transforms.RandomRotation(10), # 随机旋转
-
- transforms.ToTensor(),
-
- transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), # AlexNet的标准化
-
- ])
使用自定义数据集类和定义的数据转换来创建数据集
dataset = CustomImageDataset(main_dir="F:\\A-GX\\A-SJJ\\flower_photos\\flower_photos", transform=transform)
使用数据集创建一个数据加载器,用于批量加载和处理数据
data_loader = DataLoader(dataset, batch_size=32, shuffle=True)
从PyTorch库中加载预训练的AlexNet模型
alexnet_model = models.alexnet(pretrained=True)
修改AlexNet模型的最后几层,以便它能够处理新的分类任务
- num_ftrs = alexnet_model.classifier[6].in_features
-
- alexnet_model.classifier[6] = nn.Linear(num_ftrs, len(dataset.label_to_index))
定义用于训练模型的损失函数和优化器。
- criterion = nn.CrossEntropyLoss()
-
- optimizer = optim.Adam(alexnet_model.parameters(), lr=0.0001)
如果有多GPU,则使用nn.DataParallel来并行化模型
- if torch.cuda.device_count() > 1:
-
- alexnet_model = nn.DataParallel(alexnet_model)
将模型发送到GPU进行训练。
- device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
- alexnet_model.to(device)
数据加载器和定义的参数训练模型
- num_epochs = 10
-
- for epoch in range(num_epochs):
-
- alexnet_model.train()
-
- running_loss = 0.0
-
- for images, labels in data_loader:
-
- images, labels = images.to(device), labels.to(device)
-
-
-
- # 前向传播
-
- outputs = alexnet_model(images)
-
- loss = criterion(outputs, labels)
-
-
-
- # 反向传播和优化
-
- optimizer.zero_grad()
-
- loss.backward()
-
- optimizer.step()
-
-
-
- running_loss += loss.item()
-
-
-
- # 在每个epoch结束后评估模型
-
- train_accuracy = evaluate_model(alexnet_model, data_loader, device)
-
- print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {running_loss / len(data_loader):.4f}, Train Accuracy: {train_accuracy:.2f}%')
定义一个评估函数,用于评估模型的性能
- def evaluate_model(model, data_loader, device):
-
- model.eval() # 将模型设置为评估模式
-
- correct = 0
-
- total = 0
-
- with torch.no_grad(): # 在这个块中,所有计算都不会计算梯度
-
- for images, labels in data_loader:
-
- images, labels = images.to(device), labels.to(device)
-
- outputs = model(images)
-
- _, predicted = torch.max(outputs.data, 1)
-
- total += labels.size(0)
-
- correct += (predicted == labels).sum().item()
-
- accuracy = 100 * correct / total
-
- return accuracy