ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

PyTorch搭建CNN实现植物图像10分类实战

PyTorch搭建CNN实现植物图像10分类实战 简介这是一份基于PyTorch搭建CNN网络实现植物图像10分类的完整项目包适合正在学习深度学习图像分类的初学者和开发者可一站式解决数据集准备、模型搭建、训练与验证全流程落地问题。压缩包共198个文件以189张植物图像为主要数据来源另含4个XML标注/配置、3个Python训练脚本、1个pyc编译文件及工程配置文件整体仅10.54MB下载后即可直接运行。项目覆盖自定义Dataset读取图像、DataLoader批量加载、CNN卷积层/池化层/全连接层设计、ReLU激活、交叉熵损失训练以及验证集评估等关键环节代码注释与结构清晰便于逐模块理解和二次修改。已有711人学习下载适合作为课程设计、毕业设计或图像分类实战的参考基线也可通过调整类别数和网络层快速迁移到其他分类任务。1. 用 PyTorch 搭建 CNN 做植物图像 10 分类难点从来不在模型做植物图像分类很多人第一反应是去套用现成的 ResNet、EfficientNet结果卡在环境配置、数据目录、参数调了好久。这个项目标题里最有价值的信息不是“CNN”而是最后一句话完整代码数据可直接运行。这意味着你已经拿到了从原始图片到 10 分类预测结果的完整闭环要做的不是重新发明网络结构而是把每个环节的原理吃透再根据你自己的数据集调整参数。本文围绕 PyTorch 环境搭建、CNN 各层参数量计算、训练循环、模型保存与加载这几个核心环节展开所有代码基于 PyTorch 官方稳定版接口编写。无论你是刚接触深度学习的初学者还是已经写过一些分类模型但没跑过植物图像数据的工程师都能直接照着执行。第 2 章会先解决大多数人第一道坎环境装好了但数据集放错位置导致 ImageFolder 一个图片都读不到。2. 环境搭建与植物图像数据集组织方式2.1 用 Anaconda 创建隔离的 PyTorch 训练环境我不建议直接在 base 环境里装 PyTorch因为项目依赖的 Python 版本、CUDA 版本和 torchvision 的绑定关系很容易搞乱。用 Anaconda 单独建一个plant_cnn环境后续想删掉重来也不会影响其他项目。conda create -n plant_cnn python3.9 -y conda activate plant_cnn conda install pytorch torchvision torchaudio cpuonly -c pytorch如果你有 NVIDIA GPU 且已安装匹配的驱动把第三行替换成conda install pytorch torchvision torchaudio cudatoolkit11.8 -c pytorch -c nvidia来启用 GPU 训练。装完检查一下python -c import torch; print(torch.__version__); print(torch.cuda.is_available())输出中torch.cuda.is_available()为True才说明 GPU 可用。如果本机没有 GPUCPU 版也能训练这个量级的数据集只是单 epoch 的耗时会长一些。实际训练时建议先用 CPU 跑通全部代码确认没有逻辑错误再切换到 GPU 加速这样排错成本最低。2.2 数据集目录结构决定了 ImageFolder 能不能直接读PyTorch 的torchvision.datasets.ImageFolder依赖固定的目录层级根目录下每个子文件夹代表一个类别子文件夹里放对应类别的图片。植物图像 10 分类需要准备如下结构data/ train/ class_0/ class_1/ ... class_9/ val/ class_0/ ... class_9/训练集和验证集内部结构一致。写代码时用pathlib处理路径避免 Windows 与 Linux 路径分隔符差异导致的读取失败from pathlib import Path from torchvision import transforms, datasets from torch.utils.data import DataLoader data_root Path(data) train_data datasets.ImageFolder( rootstr(data_root / train), transformtrain_transform )ImageFolder会按照文件夹名的字母顺序自动生成类别索引train_data.classes返回的是类别名称列表train_data.class_to_idx是映射关系。这里容易踩的坑是class_10会排在class_2前面因为字符串排序不是自然排序。建议类目名都用class_0到class_9这种固定位数命名避免预测结果和真实类别对不上。2.3 数据增强策略与归一化参数设定植物图像和 ImageNet 的自然图像分布有差异但不能因此跳过归一化。标准做法是先用 ImageNet 的均值方差做 baseline后面精度不够再重新统计自己数据集的均值和方差。增强策略决定模型的泛化能力train_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(p0.5), transforms.RandomRotation(degrees15), transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.2), transforms.ToTensor(), transforms.Normalize( mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225] ) ]) val_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize( mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225] ) ]) train_loader DataLoader(train_data, batch_size32, shuffleTrue, num_workers4)RandomHorizontalFlip对植物叶片图像通常有效因为叶子朝向不影响类别判断。RandomRotation的度数范围不要超过 15 度旋转过大会把叶片形状扭曲到失去真实特征。ColorJitter四个参数前两项影响亮度对比度植物图像受光照影响大这两个值可以适当调大但饱和度和色相的调整幅度要克制否则会让模型学到错误的颜色分布。验证集不能做随机增强否则每轮验证的标准不同acc 曲线波动会很大。3. 搭建 CNN 网络模型尺寸计算与参数配置3.1 从输入到输出的张量尺寸变化规律植物图像 10 分类的输入是三通道 RGB 图像resize 到 224×224。CNN 每一层的输出特征图尺寸由公式决定output_size (input_size - kernel_size 2 * padding) / stride 1以第一层卷积为例输入 224×224×3用 3×3 卷积核、padding1、stride1输出依然是 224×224但通道数变为 32。这个设计的意义在于保持空间尺寸不缩减的前提下提取浅层特征。池化层的 kernel_size2、stride2特征图直接减半。完整网络链路中每隔一个阶段做一次下采样让网络逐步学习从边缘、纹理到器官、整体的层级特征。最后一层全局平均池化把特征图压成向量再接全连接层输出 10 个类别的 logits。3.2 用 PyTorch 定义一个可复用的植物分类 CNN 类import torch.nn as nn import torch.nn.functional as F class PlantCNN(nn.Module): def __init__(self, num_classes10): super(PlantCNN, self).__init__() self.conv_block1 nn.Sequential( nn.Conv2d(3, 32, kernel_size3, stride1, padding1), nn.BatchNorm2d(32), nn.ReLU(inplaceTrue), nn.MaxPool2d(kernel_size2, stride2) ) self.conv_block2 nn.Sequential( nn.Conv2d(32, 64, kernel_size3, stride1, padding1), nn.BatchNorm2d(64), nn.ReLU(inplaceTrue), nn.MaxPool2d(kernel_size2, stride2) ) self.conv_block3 nn.Sequential( nn.Conv2d(64, 128, kernel_size3, stride1, padding1), nn.BatchNorm2d(128), nn.ReLU(inplaceTrue), nn.MaxPool2d(kernel_size2, stride2) ) self.conv_block4 nn.Sequential( nn.Conv2d(128, 256, kernel_size3, stride1, padding1), nn.BatchNorm2d(256), nn.ReLU(inplaceTrue), nn.MaxPool2d(kernel_size2, stride2) ) self.classifier nn.Sequential( nn.Flatten(), nn.Linear(256 * 14 * 14, 512), nn.ReLU(inplaceTrue), nn.Dropout(p0.5), nn.Linear(512, num_classes) ) def forward(self, x): x self.conv_block1(x) x self.conv_block2(x) x self.conv_block3(x) x self.conv_block4(x) x self.classifier(x) return x每经过一个conv_block特征图尺寸减半224 → 112 → 56 → 28 → 14。第四个 block 输出的通道数是 256空间尺寸为 14×14展平后的向量长度为 256×14×14即 50176。这个长度直接决定了第一个全连接层的in_features参数改任何一个下采样层都会影响这个值需要同步调整。BatchNorm2d的作用是对每个通道在 batch 维度上做归一化让中间层的输入分布稳定。训练和推理模式下它的行为不同训练时用当前 batch 的统计量推理时用训练阶段累积的全局统计量。所以模型保存后在预测前必须调用model.eval()否则 BatchNorm 会因统计量不一致导致预测结果抖动。Dropout在全连接层之间随机丢弃一半神经元降低过拟合风险。这个参数在植物数据量少时很有用但如果你的训练集已经超过每类 2000 张可以适当降低到 0.3。3.3 为什么要自己搭而不直接换 ResNet自己搭的浅层 CNN 在参数量上明显少于 ResNet50更适合小数据集。ResNet50 有约 2500 万参数而上述模型只有不到 200 万在几千张植物图像上训练后者收敛更快、更难过拟合。如果你想要更强的特征提取能力一个折中方案是model torchvision.models.resnet18(pretrainedTrue) model.fc nn.Linear(512, 10)只改最后一层全连接输出为 10前面的卷积层参数保持不变。这个方案适合你已经把基础模型跑通、想要立即提升精度的阶段不建议一上来就用因为计算资源占用和训练时长会成倍增加。4. 训练循环、损失函数与优化器调参4.1 训练主循环的完整写法import torch import torch.nn as nn from torch.optim import Adam from sklearn.metrics import accuracy_score device torch.device(cuda if torch.cuda.is_available() else cpu) model PlantCNN(num_classes10).to(device) criterion nn.CrossEntropyLoss() optimizer Adam(model.parameters(), lr0.001, weight_decay5e-4) epochs 50 best_acc 0.0 for epoch in range(epochs): model.train() train_loss 0.0 all_preds [] all_labels [] for images, labels in train_loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() train_loss loss.item() * images.size(0) _, preds torch.max(outputs, 1) all_preds.extend(preds.cpu().tolist()) all_labels.extend(labels.cpu().tolist()) epoch_loss train_loss / len(train_dataset) epoch_acc accuracy_score(all_labels, all_preds) print(fEpoch {epoch1:03d} | Loss: {epoch_loss:.4f} | Acc: {epoch_acc:.4f})optimizer.zero_grad()必须在每个 batch 开始前执行清空上一轮累积的梯度否则梯度会在各 batch 间不断累加。loss.backward()计算梯度optimizer.step()用梯度更新参数。实际调试中outputs.shape是[batch_size, 10]torch.max(outputs, 1)返回的preds是每个样本的预测类别索引。Accuracy的计算可以手动写(preds labels).float().mean().item()和sklearn.metrics.accuracy_score效果一致。手动计算的好处是少引入一个依赖包。4.2 损失函数和优化器的选择依据多分类问题用CrossEntropyLoss它在 PyTorch 中已经做了 Softmax 和交叉熵的合并所以模型最后一层不需要额外加 Softmax。想输出每类的概率分布时推理阶段再加torch.softmax(outputs[0], dim0)。Adam 是自适应学习率优化器默认lr0.001在多数图像分类任务上表现稳定。weight_decay是 L2 正则化系数这里设为5e-4防止过拟合。相比 SGD 需要手动设计学习率退火策略Adam 对学习率的敏感度低适合作为 baseline。当你想追求更高精度时可以切换为optimizer torch.optim.SGD(model.parameters(), lr0.01, momentum0.9, weight_decay5e-4)SGD 配合CosineAnnealingLR学习率余弦退火往往能比 Adam 收敛到更好的局部最优点但需要更长训练时间。具体选择建议是先用 Adam 跑通流程记录最佳精度再换 SGD 精调一轮做对比。4.3 验证集评估与早停model.eval() val_loss 0.0 val_preds [] val_labels [] with torch.no_grad(): for images, labels in val_loader: images, labels images.to(device), labels.to(device) outputs model(images) loss criterion(outputs, labels) val_loss loss.item() * images.size(0) _, preds torch.max(outputs, 1) val_preds.extend(preds.cpu().tolist()) val_labels.extend(labels.cpu().tolist()) val_acc accuracy_score(val_labels, val_preds) if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), best_plant_model.pth) print(保存最佳模型)torch.no_grad()关闭梯度计算推理阶段的内存占用和计算量都会明显下降。验证集的作用不是修改模型参数而是模拟模型在未见数据上的表现。记录验证集最高的精度并保存对应权重避免训练后期过拟合导致精度回退时覆盖掉好的模型。如果训练集精度一直涨但验证集精度停滞不涨说明过拟合优先调低Dropout或者增大weight_decay。如果两边都很低说明模型容量不够或学习率不合适需要加深网络或调整lr。5. 模型保存、单张图片推理与常见环境报错5.1 保存完整 checkpoint 以便断点续训上面第 4.3 节保存的state_dict只包含模型参数不包含优化器状态。训练中断后想从中间状态继续需要保存无论数据规模如何单独的state_dict都能被load_state_dict原样加载但在代码组织上我建议把 checkpoint 与优化器状态绑定。建议保存完整 checkpointtorch.save({ epoch: epoch, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), best_acc: best_acc, }, checkpoint.tar)恢复训练时checkpoint torch.load(checkpoint.tar) model.load_state_dict(checkpoint[model_state_dict]) optimizer.load_state_dict(checkpoint[optimizer_state_dict]) start_epoch checkpoint[epoch] 1值得注意的是load_state_dict默认是严格匹配模式如果模型结构与训练时不一致例如改了全连接层的输出维度会直接报size mismatch错误。恢复训练时要保证模型定义完全一致包括num_classes参数。若模型结构做过调整加一个strictFalse参数但这样会跳过多余不匹配的键尽量不要在断点续训中用它。5.2 单张图片推理脚本import torch from PIL import Image from torchvision import transforms checkpoint torch.load(best_plant_model.pth) model PlantCNN(num_classes10) model.load_state_dict(checkpoint) model.eval() image Image.open(test_plant.jpg).convert(RGB) transform val_transform tensor transform(image).unsqueeze(0) with torch.no_grad(): outputs model(tensor) probs torch.softmax(outputs, dim1) predicted_class torch.argmax(probs, dim1).item() class_names train_data.classes print(f预测类别: {class_names[predicted_class]}, 置信度: {probs[0][predicted_class]:.4f})这段代码在保存最好是直接加载model.state_dict()来复现完整模型实例。注意加载时必须先实例化模型再load_state_dict不会直接把字典当作模型用。val_transform必须与验证集一致尤其是Resize和Normalize的参数否则输入分布变化会让预测结果明显偏差。unsqueeze(0)是在 batch 维度上增加一个维度因为单张图片的 shape 是[3, 224, 224]而模型期望[1, 3, 224, 224]。torch.argmax获取类别索引再通过class_names映射回真实类别名称。置信度可以用torch.max直接取最大值也可以用softmax后的概率值后者提供了更多参考信息如果最高置信度低于 0.5说明模型对这张图判别信心不足实际场景中应该标记为“待人工复核”。5.3 Windows 下 c10.dll 初始化失败的排查顺序在 Windows 上用 PyTorch 训练时常见报错OSError: [WinError 1114] 动态链接库(DLL)初始化例程失败。Error loading C:\Users\...\torch\lib\c10.dll。优先检查三处驱动的 CUDA 软件版本兼容性、Anaconda 环境的显卡调用依赖项、以及torch安装的版本包与 Python 版本是否匹配。conda list | findstr torch查看当前环境的 PyTorch 相关包版本确认安装的是cpuonly或是cpu构建还是cuda构建。如果本机没有 NVIDIA 显卡但装的是 CUDA 构建卸载后重装cpuonly版本即能解决。有 NVIDIA 显卡的朋友需要确保nvidia-smi输出的驱动版本与 CUDA 版本兼容。基于 PyTorch 的 CNN 网络实现植物图像分类完成数据准备、模型构建、训练迭代、验证评估这套流程后可以继续尝试把输入分辨率从 224 降到 96 来加速实验迭代或者用torch.optim.lr_scheduler.StepLR做阶梯式学习率衰减进一步的性能优化方向可以落在数据增强的自动化搜索和特征图的 Grad-CAM 可视化上。本文还有配套的精品资源点击获取
返回列表