当前位置：首页 > news >正文

从零开始：用PyTorch和Swin Transformer搞定花卉图像分类（附完整代码和常见报错解决）

news 2026/6/9 0:05:22

从零开始：用PyTorch和Swin Transformer搞定花卉图像分类（附完整代码和常见报错解决）

当你在花园里看到一朵花，能立刻叫出它的名字吗？对于大多数人来说，识别花卉种类可能是个挑战，但对计算机来说，这正是一个典型的图像分类问题。本文将带你从零开始，使用PyTorch和最新的Swin Transformer模型，构建一个强大的花卉分类系统。

1. 环境准备与工具选择

在开始之前，我们需要搭建一个合适的工作环境。以下是推荐的配置：

Python 3.8+：确保使用较新的Python版本
PyTorch 1.7+：深度学习框架的基础
CUDA 11.0+（如果使用GPU加速）
timm库：包含预训练的Swin Transformer模型

安装核心依赖的命令如下：

pip install torch torchvision timm matplotlib opencv-python

为什么选择Swin Transformer而不是传统的CNN模型？这里有一个简单的对比：

特性	CNN模型 (如ResNet)	Swin Transformer
局部特征提取能力	强	中等
全局上下文理解能力	弱	强
计算效率	高	中等
数据需求	相对较少	较多
可解释性	较好	较差

对于花卉分类这种需要理解整体形态和局部细节的任务，Swin Transformer的层次化窗口注意力机制表现出色。

2. 数据集准备与预处理

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

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

数据预处理是关键步骤，以下是推荐的转换流程：

from torchvision import transforms train_transform = transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) val_transform = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ])

常见的数据处理问题及解决方案：

类别不平衡：某些花卉的样本可能比其他类别多
- 解决方案：使用加权采样或数据增强
图像尺寸不一：花朵在图片中的比例不同
- 解决方案：统一缩放到224x224（Swin-T的默认输入尺寸）
过拟合风险：花卉数据集通常样本量有限
- 解决方案：使用迁移学习和数据增强

3. 模型构建与训练

Swin Transformer的核心优势在于其层次化的窗口注意力机制，能够高效处理视觉任务。以下是构建模型的代码示例：

import torch import timm def build_model(num_classes=5): model = timm.create_model('swin_tiny_patch4_window7_224', pretrained=True) model.head = torch.nn.Linear(model.head.in_features, num_classes) return model

训练过程中需要注意的关键点：

学习率设置：Swin Transformer通常需要较小的学习率
优化器选择：AdamW通常表现良好
训练策略：使用学习率warmup和余弦衰减

训练代码框架：

model = build_model().to(device) criterion = torch.nn.CrossEntropyLoss() optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4) for epoch in range(epochs): model.train() for images, labels in train_loader: images, labels = images.to(device), labels.to(device) outputs = model(images) loss = criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step()

4. 常见问题与解决方案

在实际项目中，你可能会遇到以下典型问题：

4.1 _IncompatibleKeys警告

这是预训练模型加载时常见的问题，表示部分权重无法匹配。解决方案：

# 加载预训练权重时忽略不匹配的键 model.load_state_dict(torch.load(weights_path), strict=False)

4.2 KeyboardInterrupt错误

在predict.py运行时出现的键盘中断通常是由于matplotlib的显示问题。解决方案：

# 修改predict.py中的显示代码 plt.show(block=True) # 添加block参数

4.3 内存不足问题

Swin Transformer虽然高效，但仍需要较多显存。解决方法包括：

减小batch size
使用混合精度训练
尝试更小的模型变体

4.4 过拟合问题

当验证集准确率远低于训练集时，可以尝试：

增加数据增强
添加Dropout层
使用更早的停止策略
减少模型复杂度

5. 模型评估与部署

训练完成后，我们需要评估模型性能。关键指标包括：

整体准确率
各类别的精确率、召回率
混淆矩阵

部署模型进行预测的示例代码：

def predict(image_path, model, transform): img = Image.open(image_path).convert('RGB') img_t = transform(img).unsqueeze(0) with torch.no_grad(): outputs = model(img_t) probs = torch.nn.functional.softmax(outputs, dim=1) return probs.squeeze().tolist()

为了提高预测效率，可以考虑以下优化：

使用ONNX格式导出模型
实现批处理预测
使用TensorRT加速

6. 进阶技巧与优化

要让你的花卉分类器更上一层楼，可以尝试以下技巧：

自定义数据增强：
- 针对花卉特点设计旋转、颜色抖动等增强
- 使用Albumentations库实现复杂增强
模型微调策略：
- 分层设置学习率（后面层的学习率更高）
- 逐步解冻模型层
集成学习：
- 结合多个Swin Transformer模型的预测结果
- 与CNN模型集成
注意力可视化：
- 可视化Swin Transformer的注意力图
- 分析模型关注的花卉关键特征

# 注意力可视化示例代码 def visualize_attention(image, model): # 注册hook获取注意力图 attentions = [] def hook(module, input, output): attentions.append(output[1]) # 获取注意力权重 handle = model.blocks[0].attn.register_forward_hook(hook) # 前向传播 model(image.unsqueeze(0)) # 移除hook handle.remove() return attentions[0]