ARTICLE DETAIL

资讯详情

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

基于CNN的大米品种识别实战:从数据集预处理到模型部署全流程

基于CNN的大米品种识别实战:从数据集预处理到模型部署全流程 简介这份资源面向深度学习入门者与计算机视觉方向的开发者提供一套基于PyTorch框架的CNN大米识别完整实现方案可用于图像分类实战练习或课程设计参考。压缩包共906个文件包含900张jpg图片构成的多类别数据集、3个txt说明与日志文本以及3个py脚本整体约11.98MB体积轻便便于本地运行。数据集已做预处理通过短边补灰边将图片统一为正方形并辅以旋转、翻转等方式扩增样本提升模型泛化能力。代码流程清晰先运行数据集文本生成脚本读取各类别图片路径与标签再执行模型训练脚本完成训练并保存本地模型同时输出记录每个epoch验证集损失与准确率的日志最后通过PyQt界面加载图片进行可视化识别。已有152人学习适合希望打通数据预处理、模型训练到界面部署全链路的读者参考。1. 大米识别为什么值得用 CNN 做一遍从一张米粒图说起大米品种识别这件事听起来像是农业质检的小众需求但真正做过的人都知道它是个典型的「看着简单、做起来全是坑」的视觉分类任务。你拿一张米粒照片给普通人看他能分出长粒香和珍珠米已经算厉害了可要是让你区分稻花香、中早39、黄华占这些品种别说肉眼就是做了十几年粮食收购的老师傅也得靠经验加运气。而基于 CNN 深度学习的大米识别本质上就是让卷积神经网络去学那些人类说不清、但确实存在的纹理和形状差异。这个方向适合谁一是做农产品质检自动化的工程师二是想拿一个完整数据集练手深度学习的新手三是做粮食仓储、加工分选设备的技术负责人。它不需要 GPU 集群一台带独显的笔记本就能跑通全流程但它也不是「随便调个库就完事」的玩具项目数据集的图片质量、预处理方式、网络结构选择每一步都会直接影响最终精度。我见过太多人拿现成的 CNN 模板套上去结果测试集准确率卡在 70% 上不去最后发现是米粒背景没统一、光照差异太大导致的。所以这篇笔记我按「数据集怎么用 → 模型怎么搭 → 训练怎么调 → 坑怎么避」的顺序把这条链路完整走一遍。2. 图片数据集到手后先别急着训练拆解与预处理2.1 数据集结构长什么样先摸清楚再动手拿到「含图片数据集.zip」这种资源第一件事不是解压完就往模型里灌而是先搞清楚它的目录结构。常见的大米图片数据集一般按品种分文件夹每个文件夹下是该品种的米粒照片可能是单粒拍摄也可能是一小堆。你需要确认三件事类别数、每类图片数量、图片尺寸是否统一。我一般会先跑一段脚本做统计而不是靠肉眼翻文件夹。下面这段代码就是干这个的import os from PIL import Image from collections import defaultdict data_dir ./rice_dataset # 解压后的根目录 stats defaultdict(int) size_set set() for class_name in os.listdir(data_dir): class_path os.path.join(data_dir, class_name) if not os.path.isdir(class_path): continue for img_file in os.listdir(class_path): img_path os.path.join(class_path, img_file) try: with Image.open(img_path) as im: size_set.add(im.size) # 记录所有出现过的尺寸 stats[class_name] 1 except Exception as e: print(f坏图: {img_path}, 原因: {e}) print(类别分布:, dict(stats)) print(出现过的图片尺寸:, size_set) print(总类别数:, len(stats))这段逻辑很直白遍历每个类别文件夹统计图片数量同时把所有图片的尺寸收集到一个集合里。如果size_set里出现超过三种尺寸说明数据集里的图片没有统一过后面必须做 resize如果某个类别只有几十张而其他类别有几百张那就是典型的类别不平衡训练时得加权重或者做增强。参数上没什么好调的重点看输出结果——类别分布是否均匀、尺寸是否一致这两个信息决定了你后面预处理管线的复杂度。2.2 预处理管线的四个关键参数怎么定确认完数据集的基本情况接下来就是搭预处理管线。用 PyTorch 的torchvision.transforms也好用 TensorFlow 的ImageDataGenerator也好核心参数就四个resize 尺寸、归一化均值方差、增强策略、验证集划分比例。resize 尺寸我一般定在 224×224 或 256×256。224 是 ImageNet 系列模型的标准输入用预训练权重时不用改网络第一层256 则保留更多纹理细节适合米粒这种靠表面纹路区分的任务。如果你的数据集图片本身就不大比如 300×300 左右直接 resize 到 224 损失的信息有限可以接受。归一化均值和方差如果是从零训练用数据集自己算出来的均值和方差最合适如果用预训练模型就直接用 ImageNet 的mean[0.485, 0.456, 0.406]、std[0.229, 0.224, 0.225]。这一点很多人会忽略觉得归一化随便填就行但实际上均值方差不对预训练权重的特征分布就对不上收敛会慢很多。增强策略方面大米识别任务里水平翻转、小角度旋转±15°、亮度微调是安全且有效的但垂直翻转要慎用因为米粒的胚芽位置有方向性垂直翻转可能产生现实中不存在的样本。验证集划分比例一般取 15%20%如果每类图片少于 200 张建议用 5 折交叉验证而不是单次划分。from torchvision import transforms train_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(p0.5), transforms.RandomRotation(degrees15), transforms.ColorJitter(brightness0.2, contrast0.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]) ])注意验证集的 transform 里不能加随机增强只能做 resize 和归一化否则验证指标会抖动你没法判断模型是真的进步了还是增强带来的随机性。这个坑我踩过不止一次血泪经验就是训练管线和验证管线一定要分开写别图省事共用一个 Compose。3. 用 CNN 搭一个能跑通的大米分类器从 LeNet 到 ResNet 的选型3.1 网络结构选型别一上来就 ResNet-152大米识别这个任务的复杂度介于 MNIST 和 ImageNet 之间。类别数通常不超过 20 类图片内容单一就是米粒背景相对可控。这种任务用太深的网络参数量远大于实际需求训练慢还容易过拟合用太浅的网络比如 LeNet-5特征提取能力又不够精度上不去。我的建议是分三档如果类别数少于 10 类、每类图片超过 500 张用 ResNet-18 或 MobileNetV2 就够了如果类别数在 1030 类之间用 ResNet-34 或 EfficientNet-B0如果类别数超过 30 类或者图片背景很复杂再考虑 ResNet-50 起步。预训练权重能加载就加载哪怕你的数据集和 ImageNet 差异大底层卷积核学到的边缘、纹理特征也是通用的比随机初始化收敛快得多。下面是一个基于 ResNet-18 的完整模型定义包括替换分类头和冻结策略import torch import torch.nn as nn from torchvision import models def build_model(num_classes, freeze_backboneTrue): model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) if freeze_backbone: for param in model.parameters(): param.requires_grad False # 替换最后的全连接层适配大米类别数 in_features model.fc.in_features model.fc nn.Sequential( nn.Dropout(0.3), nn.Linear(in_features, num_classes) ) return model model build_model(num_classes12, freeze_backboneTrue)这里freeze_backboneTrue表示冻结主干网络只训练新加的分类头。这样做的好处是训练快、不容易过拟合适合数据集规模不大的情况。等分类头收敛后可以解冻最后几个 block 做微调学习率调小一个数量级。Dropout(0.3)是我习惯加的正则化大米识别任务里过拟合比欠拟合常见得多。3.2 训练循环里必须监控的三个量训练循环本身没什么特别的但有三个量你必须盯着训练损失、验证损失、验证准确率。训练损失下降但验证损失上升说明过拟合了两个都下降但验证准确率不动说明学习率可能太小训练损失震荡剧烈说明 batch size 太小或者学习率太大。import torch.optim as optim from torch.utils.data import DataLoader criterion nn.CrossEntropyLoss() optimizer optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr1e-3) scheduler optim.lr_scheduler.StepLR(optimizer, step_size10, gamma0.1) for epoch in range(30): model.train() running_loss 0.0 for images, labels in train_loader: optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() scheduler.step() # 验证阶段 model.eval() correct, total 0, 0 val_loss 0.0 with torch.no_grad(): for images, labels in val_loader: outputs model(images) val_loss criterion(outputs, labels).item() _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() print(fEpoch {epoch1}: train_loss{running_loss/len(train_loader):.4f}, fval_loss{val_loss/len(val_loader):.4f}, val_acc{correct/total:.4f})学习率我一般从 1e-3 开始配合 StepLR 每 10 个 epoch 降一次。如果验证准确率在 5 个 epoch 内没有提升就把学习率降到 1e-4 再试。batch size 在显存允许的前提下尽量大16 或 32 都行太小了梯度噪声大训练曲线会很难看。3.3 类别不平衡时用加权损失还是重采样如果第 2 章统计出来发现类别不平衡比如最多的类有 800 张、最少的只有 120 张直接训练会导致模型偏向多数类。两种处理方式一是给 CrossEntropyLoss 加weight参数少数类权重调高二是用 WeightedRandomSampler 做重采样让每个 batch 里各类别比例均衡。from torch.utils.data import WeightedRandomSampler import numpy as np class_counts [800, 650, 400, 120] # 按类别顺序 class_weights 1.0 / torch.tensor(class_counts, dtypetorch.float) sample_weights [class_weights[label] for _, label in train_dataset] sampler WeightedRandomSampler(sample_weights, num_sampleslen(sample_weights), replacementTrue) train_loader DataLoader(train_dataset, batch_size32, samplersampler)重采样的好处是每个 epoch 都能见到少数类的样本缺点是少数类可能被重复采样导致过拟合。我的经验是不平衡比例在 3:1 以内用加权损失就够了超过 5:1重采样效果更明显。两者也可以叠加使用但要注意别把少数类权重调得过高否则模型会对少数类过拟合验证集上反而掉点。4. 训练完不算完评估、调参和踩坑记录4.1 混淆矩阵比准确率更能说明问题准确率这个指标在大米识别任务里很容易骗人。如果 12 个类别里有 10 个类别的准确率都在 95% 以上但有两个类别互相混淆严重整体准确率可能还有 90%看起来不错实际上那两个类别根本没法用。所以训练完之后一定要画混淆矩阵看看错误到底集中在哪些类别对之间。from sklearn.metrics import confusion_matrix, classification_report import seaborn as sns import matplotlib.pyplot as plt model.eval() all_preds, all_labels [], [] with torch.no_grad(): for images, labels in val_loader: outputs model(images) _, predicted torch.max(outputs, 1) all_preds.extend(predicted.cpu().numpy()) all_labels.extend(labels.cpu().numpy()) cm confusion_matrix(all_labels, all_preds) plt.figure(figsize(10, 8)) sns.heatmap(cm, annotTrue, fmtd, cmapBlues) plt.xlabel(Predicted) plt.ylabel(True) plt.show() print(classification_report(all_labels, all_preds, target_namesclass_names))如果发现某两个类别互相混淆严重先别急着调网络结构去看看这两个类别的图片是不是本身就很像——比如都是长粒型、颜色也接近。这种情况下要么补充更多区分性强的样本要么在预处理阶段加入更强的纹理增强比如随机锐化、边缘增强让模型学到更细粒度的差异。4.2 避坑记录大米识别项目里最常见的五个翻车现场现象一训练集准确率 99%验证集只有 65%。原因典型过拟合数据集太小或者增强不够。解决先冻结主干只训分类头再加 Dropout 和权重衰减如果还不行就扩充数据集或改用更强的增强策略。现象二训练损失从第一个 epoch 就不下降。原因学习率太大或者归一化参数不对。解决把学习率降到 1e-4 试试同时检查归一化均值和方差是否和预训练模型匹配。如果是从零训练确认输入图片的像素值是否已经缩放到 [0,1] 区间。现象三验证准确率波动很大每次跑结果都不一样。原因验证集太小或者验证集用了随机增强。解决扩大验证集比例到 20%确保验证 transform 里没有随机操作。如果数据集本身小改用交叉验证。现象四某些类别始终识别不出来。原因类别不平衡或者这些类别的图片质量差。解决先检查这些类别的图片是不是有模糊、遮挡、光照异常的问题清理坏图后再用重采样训练。现象五模型在测试集上表现好但实际拍照识别就翻车。原因训练数据和实际场景的域差异太大。解决在预处理里加入更贴近实际场景的增强比如随机背景替换、不同光照模拟、运动模糊等。如果条件允许补充实际场景下拍摄的图片做微调。5. 把模型推到能用的程度几个我反复验证过的技巧5.1 用测试时增强把精度再挤两个点模型训练完之后如果精度还差一点到你的要求可以试试测试时增强TTA。做法很简单对同一张测试图片做多次不同的变换比如原图、水平翻转、小角度旋转分别推理后把 softmax 输出平均取最大概率对应的类别。这个技巧几乎不增加训练成本但通常能带来 13 个点的精度提升。def predict_with_tta(model, image, n_aug5): model.eval() probs torch.zeros(1, num_classes) with torch.no_grad(): for _ in range(n_aug): aug_img val_transform(image) # 实际使用时可以加随机增强 output model(aug_img.unsqueeze(0)) probs torch.softmax(output, dim1) return probs.argmax(dim1).item()注意 TTA 只适合验证和推理阶段训练时不能用否则验证指标会虚高。另外增强的次数不是越多越好5 到 8 次就够了再多边际收益很低。5.2 模型导出和推理速度优化如果这个大米识别模型最终要部署到分选设备上推理速度就是个硬指标。ResNet-18 在 CPU 上单张推理大概 3050msGPU 上 5ms 以内。如果设备只有 CPU 而且要求实时可以考虑把模型转成 ONNX 或者用 TorchScript 优化再配合量化把 FP32 转成 INT8速度能提升 23 倍精度损失通常在 1 个点以内。# 导出为 ONNX dummy_input torch.randn(1, 3, 224, 224) torch.onnx.export(model, dummy_input, rice_cnn.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch}, output: {0: batch}})导出之后记得用 onnxruntime 跑一遍验证确认输出和 PyTorch 一致。我一般会保留 PyTorch 版本作为基准ONNX 版本作为部署版本两边同时跑一批测试图片对比结果差异。5.3 我踩过的最大的一个坑最后说一个我自己的教训。有一次我拿到一个标注好的大米数据集训练集准确率冲到 98%验证集也有 95%我兴冲冲地拿去实际场景测试结果直接掉到 60%。排查了一整天最后发现数据集里的图片全是在固定光源、固定背景下拍的而实际场景的光源色温和背景颜色完全不同。模型学到的是「背景颜色 米粒形状」的联合特征而不是纯粹的米粒纹理。从那以后我养成了一个习惯拿到任何图像数据集先花半小时翻一遍图片看看拍摄条件是否统一、有没有明显的域偏移。如果训练数据和实际场景差异大宁可在预处理阶段多花时间做域增强也不要等模型训完了再返工。这个习惯帮我省下的时间远比那半小时多得多。希望帮到你。本文还有配套的精品资源点击获取
返回列表