AlexNet实战：用PyTorch从零搭建花卉分类模型（附完整代码+数据集）

AlexNet实战：用PyTorch从零搭建花卉分类模型

深度学习在计算机视觉领域的应用已经变得无处不在，而图像分类作为最基础的任务之一，仍然是初学者入门的最佳选择。本文将带你从零开始，使用PyTorch框架实现经典的AlexNet模型，并应用于花卉分类任务。不同于简单的API调用，我们会深入模型架构的每个细节，让你真正理解卷积神经网络的工作原理。

1. 环境准备与数据预处理

在开始构建模型之前，我们需要准备好开发环境和数据集。PyTorch作为当前最流行的深度学习框架之一，以其动态计算图和Pythonic的API设计赢得了大量开发者的青睐。

1.1 安装必要的依赖

首先确保你已经安装了Python 3.7或更高版本，然后通过pip安装以下包：

pip install torch torchvision pillow matplotlib numpy tqdm

对于GPU加速，还需要安装对应版本的CUDA和cuDNN。可以使用以下命令检查PyTorch是否正确识别了你的GPU：

import torch print(torch.cuda.is_available()) # 应该输出True如果有可用的GPU

1.2 准备花卉数据集

我们将使用一个包含5类花卉的公开数据集，类别包括：

• 雏菊(daisy)
• 蒲公英(dandelion)
• 玫瑰(roses)
• 向日葵(sunflowers)
• 郁金香(tulips)

数据集的组织结构应该如下：

flower_data/ train/ daisy/ image1.jpg image2.jpg ... dandelion/ roses/ sunflowers/ tulips/ val/ daisy/ dandelion/ roses/ sunflowers/ tulips/

提示：可以使用split_data.py脚本自动划分训练集和验证集，确保验证集约占全部数据的10%-20%。

1.3 数据增强与预处理

在深度学习中，数据预处理是至关重要的一步。对于图像分类任务，我们通常需要进行以下操作：

from torchvision import transforms # 训练集的数据增强 train_transform = transforms.Compose([ transforms.RandomResizedCrop(224), # 随机裁剪并缩放到224x224 transforms.RandomHorizontalFlip(), # 随机水平翻转 transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2), # 颜色抖动 transforms.ToTensor(), # 转换为Tensor并归一化到[0,1] transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # ImageNet标准化 ]) # 验证集的预处理（不需要数据增强） val_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]) ])

数据增强可以有效防止过拟合，特别是在数据集较小的情况下。通过随机变换输入图像，我们实际上是在"创造"新的训练样本。

2. AlexNet模型架构详解

AlexNet作为深度学习的里程碑，在2012年ImageNet竞赛中以显著优势夺冠，开启了深度学习在计算机视觉领域的新时代。虽然现在有更先进的模型，但理解AlexNet仍然是学习CNN的重要一步。

2.1 网络结构分析

AlexNet的原始架构包含5个卷积层和3个全连接层，由于当时GPU内存限制，设计为在两个GPU上并行计算。在我们的实现中，我们将所有参数减半以适应现代消费级GPU。

完整的AlexNet架构如下表所示：

2.2 PyTorch实现

下面是完整的AlexNet实现代码：

import torch.nn as nn import torch class AlexNet(nn.Module): def __init__(self, num_classes=5, init_weights=True): super(AlexNet, self).__init__() self.features = nn.Sequential( nn.Conv2d(3, 48, kernel_size=11, stride=4, padding=2), # [3,224,224]→[48,55,55] nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=3, stride=2), # [48,55,55]→[48,27,27] nn.Conv2d(48, 128, kernel_size=5, padding=2), # [48,27,27]→[128,27,27] nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=3, stride=2), # [128,27,27]→[128,13,13] nn.Conv2d(128, 192, kernel_size=3, padding=1), # [128,13,13]→[192,13,13] nn.ReLU(inplace=True), nn.Conv2d(192, 192, kernel_size=3, padding=1), # [192,13,13]→[192,13,13] nn.ReLU(inplace=True), nn.Conv2d(192, 128, kernel_size=3, padding=1), # [192,13,13]→[128,13,13] nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=3, stride=2), # [128,13,13]→[128,6,6] ) self.classifier = nn.Sequential( nn.Dropout(p=0.5), nn.Linear(128 * 6 * 6, 2048), nn.ReLU(inplace=True), nn.Dropout(p=0.5), nn.Linear(2048, 2048), nn.ReLU(inplace=True), nn.Linear(2048, num_classes), ) if init_weights: self._initialize_weights() def forward(self, x): x = self.features(x) x = torch.flatten(x, start_dim=1) x = self.classifier(x) return x def _initialize_weights(self): for m in self.modules(): if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') if m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.Linear): nn.init.normal_(m.weight, 0, 0.01) nn.init.constant_(m.bias, 0)

关键点说明：

1. inplace=True的ReLU可以节省一些内存
2. Kaiming初始化特别适合ReLU激活函数
3. Dropout层只在训练时激活，可以防止过拟合
4. 最后一层不需要激活函数，因为我们将使用CrossEntropyLoss

3. 模型训练与调优

有了模型架构后，我们需要设置训练流程。这部分将详细介绍如何高效地训练AlexNet模型。

3.1 训练配置

首先设置训练参数和优化器：

import torch.optim as optim from torch.utils.data import DataLoader # 初始化模型 model = AlexNet(num_classes=5, init_weights=True) model = model.to(device) # 移动到GPU如果可用 # 定义损失函数和优化器 criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=0.0002) # 学习率调度器 scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)

注意：Adam优化器通常比原始论文中使用的SGD with momentum表现更好，特别是对于初学者。

3.2 训练循环

完整的训练循环包括以下几个步骤：

1. 前向传播计算输出
2. 计算损失
3. 反向传播计算梯度
4. 优化器更新权重
5. 定期在验证集上评估模型

def train_model(model, criterion, optimizer, scheduler, num_epochs=10): best_acc = 0.0 for epoch in range(num_epochs): print(f'Epoch {epoch}/{num_epochs - 1}') print('-' * 10) # 每个epoch都有训练和验证阶段 for phase in ['train', 'val']: if phase == 'train': model.train() # 设置模型为训练模式 else: model.eval() # 设置模型为评估模式 running_loss = 0.0 running_corrects = 0 # 迭代数据 for inputs, labels in dataloaders[phase]: inputs = inputs.to(device) labels = labels.to(device) # 梯度清零 optimizer.zero_grad() # 前向传播 with torch.set_grad_enabled(phase == 'train'): outputs = model(inputs) _, preds = torch.max(outputs, 1) loss = criterion(outputs, labels) # 只在训练阶段反向传播+优化 if phase == 'train': loss.backward() optimizer.step() # 统计 running_loss += loss.item() * inputs.size(0) running_corrects += torch.sum(preds == labels.data) if phase == 'train': scheduler.step() epoch_loss = running_loss / dataset_sizes[phase] epoch_acc = running_corrects.double() / dataset_sizes[phase] print(f'{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}') # 深度复制模型 if phase == 'val' and epoch_acc > best_acc: best_acc = epoch_acc torch.save(model.state_dict(), 'best_model.pth') print() print(f'Best val Acc: {best_acc:4f}') return model

3.3 训练技巧与调优

在实际训练中，有几个关键技巧可以提升模型性能：

1. 学习率调整：初始学习率设置为0.0002，每5个epoch乘以0.1
2. 早停(Early Stopping)：如果验证集准确率连续几个epoch不提升，可以提前终止训练
3. 模型检查点：保存验证集上表现最好的模型
4. 梯度裁剪：防止梯度爆炸

# 添加梯度裁剪 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) # 早停实现 patience = 3 # 允许的连续不提升epoch数 no_improve = 0 best_loss = float('inf') for epoch in range(num_epochs): # ...训练代码... val_loss = epoch_loss if val_loss < best_loss: best_loss = val_loss no_improve = 0 torch.save(model.state_dict(), 'best_model.pth') else: no_improve += 1 if no_improve >= patience: print("Early stopping!") break

4. 模型评估与预测

训练完成后，我们需要评估模型性能并进行实际预测。

4.1 评估模型性能

使用混淆矩阵可以直观地展示模型在各个类别上的表现：

from sklearn.metrics import confusion_matrix import seaborn as sns import pandas as pd def plot_confusion_matrix(model, dataloader, class_names): model.eval() all_preds = [] all_labels = [] with torch.no_grad(): for inputs, labels in dataloader: inputs = inputs.to(device) labels = labels.to(device) outputs = model(inputs) _, preds = torch.max(outputs, 1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.cpu().numpy()) cm = confusion_matrix(all_labels, all_preds) df_cm = pd.DataFrame(cm, index=class_names, columns=class_names) plt.figure(figsize=(10,7)) sns.heatmap(df_cm, annot=True, fmt='d', cmap='Blues') plt.xlabel('Predicted') plt.ylabel('True') plt.show()

4.2 单张图像预测

下面是一个完整的预测流程，可以用于单张图像的分类：

def predict_image(image_path, model, transform, class_names): # 加载图像 img = Image.open(image_path) # 预处理 img_t = transform(img) batch_t = torch.unsqueeze(img_t, 0).to(device) # 预测 model.eval() with torch.no_grad(): output = model(batch_t) # 获取预测结果 _, pred = torch.max(output, 1) prob = torch.nn.functional.softmax(output, dim=1)[0] * 100 # 显示结果 plt.imshow(img) plt.title(f'Predicted: {class_names[pred.item()]} ({prob[pred.item()]:.1f}%)') plt.axis('off') plt.show() # 打印所有类别概率 for i, (name, p) in enumerate(zip(class_names, prob)): print(f'{name}: {p:.1f}%') return class_names[pred.item()]

4.3 可视化中间特征

理解CNN工作原理的一个好方法是可视化中间层的特征图：

def visualize_feature_maps(model, image_path, layer_index=0): # 加载并预处理图像 img = Image.open(image_path) 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]) ]) img_t = transform(img).unsqueeze(0).to(device) # 获取指定层的输出 activation = {} def get_activation(name): def hook(model, input, output): activation[name] = output.detach() return hook target_layer = list(model.features.children())[layer_index] handle = target_layer.register_forward_hook(get_activation(f'conv{layer_index+1}')) # 前向传播 with torch.no_grad(): output = model(img_t) # 可视化特征图 act = activation[f'conv{layer_index+1}'].squeeze().cpu() fig, axarr = plt.subplots(act.size(0)//8, 8, figsize=(20, 20)) for idx in range(act.size(0)): ax = axarr[idx//8, idx%8] ax.imshow(act[idx], cmap='viridis') ax.axis('off') plt.tight_layout() plt.show() # 移除hook handle.remove()

这个可视化可以帮助我们理解CNN每一层学习到了什么样的特征。通常，浅层会学习边缘、颜色等低级特征，而深层会学习更抽象的特征。