0 前言：杂草识别为什么难？

在智慧农业落地场景中，杂草识别是精准除草机器人的核心前置环节。现实农田环境充满挑战：不同杂草幼苗外观高度相似；光照变化、泥土背景、叶片遮挡都会干扰模型判断；传统纯卷积CNN依靠局部卷积核提取特征，感受野有限，很难捕捉杂草整体形态差异。

普通ResNet虽然依靠残差连接解决深度网络梯度消失，在通用图像分类任务表现优秀，但卷积运算天然偏向局部信息，对细粒度、类间相似度极高的杂草数据集，容易出现混淆错分。Transformer的自注意力擅长全局依赖建模，但参数量大、训练开销高，直接替换主干成本太高。

因此本项目采用CNN+自注意力混合架构：以ResNet34作为卷积主干，在4个残差阶段后分别嵌入2D自注意力SelfAttention模块。既保留卷积强大局部纹理提取能力，又引入全局注意力机制，低成本提升细粒度杂草识别效果。项目使用4700+张田间杂草图片，完成15种杂草分类，代码包含完整训练train.py、批量预测predict.py、工具集utils.py，自动输出全部实验指标与可视化图表。

数据集包含15类杂草：Carpetweeds、Crabgrass、Eclipta、Goosegrass、Morningglory、Nutsedge、PalmerAmaranth、Prickly Sida、Purslane、Ragweed、Sicklepod、SpottedSpurge、SpurredAnoda、Swinecress、Waterhemp；类别索引与名称映射保存class_indices.json，训练推理统一读取该标签文件。

1 整体网络架构详解

整体流程：输入杂草图像→图像预处理（Resize、随机旋转、中心裁剪、归一化）→ResNet34主干网络（Conv1+MaxPool + layer1~layer4残差层）→每一个layer后插入SelfAttention自注意力模块→全局平均池化→全连接FC输出15分类Logits。

Mermaid网络结构图：

1.1 SelfAttention自注意力模块实现原理

项目utils.py中实现二维空间自注意力模块，不使用多头，通过1×1卷积生成Query、Key、Value：

1. 输入特征图[B,C,H,W]；
2. 使用1×1卷积分别得到Q、K、V，通道降维至C//8减少计算；
3. 将Q、K变形，矩阵乘法计算空间注意力矩阵；
4. Softmax得到空间注意力权重，对Value特征加权聚合；
5. 学习系数gamma做残差加权，out = gamma * attention_feature + input，避免直接引入注意力破坏原有卷积特征分布。

关键设计：在layer1‑layer4全部残差阶段之后插入SA模块。浅层注意力捕捉叶片边缘、纹理细节；深层注意力建模杂草整体植株空间分布，强化区分形态近似杂草。

1.2 ResNet34主干与迁移学习策略

主干采用torchvision预训练ResNet34，加载ImageNet预训练权重做迁移学习。代码支持freeze_layers开关：开启冻结时只训练FC分类头和所有SelfAttention模块，主干卷积参数固定；关闭冻结时全部网络参数参与更新。这种策略在数据集只有4700张中等规模数据场景下，可以有效防止过拟合，加快收敛速度。

2 数据集、数据增强与实验配置

• 数据集目录结构：data/train、data/val，使用ImageFolder自动加载，自动生成class_indices.json类别映射文件；
• 训练集增强：Resize256，随机旋转90°，CenterCrop224，ToTensor，ImageNet均值方差归一化；
• 验证集：无随机旋转，仅做缩放裁剪归一化；
• 训练超参数：batch‑size=8，epochs=30；优化器可选SGD / Adam；基础学习率lr=0.001；采用cosine余弦退火学习率调度器；
• 损失函数：交叉熵CrossEntropyLoss。

3 完整训练流程（train.py逻辑梳理）

1. 初始化运行目录runs，创建权重保存文件夹；获取设备自动选择CUDA/GPU或CPU；
2. 加载训练集、验证集，打印样本数量，生成类别json文件；
3. 构建改进ResNet34‑SA模型，打印参数量、FLOPs；区分可训练参数；
4. 实例化优化器、余弦学习率scheduler；
5. 训练循环每个epoch：
   • train_one_epoch：前向、损失计算、反向更新梯度，更新训练混淆矩阵，返回训练loss、lr、各项指标；
   • scheduler.step更新学习率；
   • evaluate函数跑验证集，计算val loss与验证混淆矩阵；
   • 保存每一轮全部指标到json日志；
   • 当验证集准确率创新高，保存best.pth；训练结束保存last.pth；
6. 训练结束后自动可视化：
   • loss‑accuracy曲线；
   • 学习率衰减曲线；
   • Precision、Recall、Specificity、F1四个指标随epoch变化曲线plot_prsf；
   • 绘制训练集、验证集混淆矩阵图片；
7. 混淆矩阵模块ConfusionMatrix可以输出每一类杂草的精确率、召回率、特异度、F1，同时输出各类指标均值，方便评估细粒度分类效果。

坑点说明：在evaluate函数源码中存在一处小隐患，循环内部写了model.train()，评估阶段应该设置model.eval()，会影响BN、Dropout运行状态，实际使用建议修正。

4 批量推理模块predict.py详解

推理文件夹结构inference/infer_img存放待预测图片；推理后结果输出到inference/ret。

1. 读取class_indices.json类别映射；
2. 加载best.pth权重，模型设置eval模式，torch.no_grad()关闭梯度；
3. 遍历文件夹下全部图片，PIL读取，做和验证集一致的transform；
4. 网络前向+softmax得到各类别概率；torch.topk取Top‑3（类别>2）或Top‑2预测结果；
5. OpenCV把类别名称+置信度绘制到原图上，保存结果图片。

适合直接部署在本地做杂草图片批量分析，输出可视化结果，直观观察模型判断结果。

融合自注意力的 ResNet34 农田杂草细粒度分类研究