PyTorch + EfficientNet 训练农作物病害检测模型_农作物病害数据训练-CSDN博客

本文链接：https://blog.csdn.net/weixin_41323338/article/details/146354682

病虫害是农作物生产中的重大挑战，准确检测病害有助于快速准确的帮助专家和农民及时防治。我这次使用了一个小麦的病害的数据集来测试训练模型。

首先需要数据增强，如果你数据集较小，更需要数据增强，以防止过拟合。通常需要上千张的图片来确定一个病害的外表特征。因为病害特征比较多变，这就需要不同的光亮，角度等不同情况下的图片。而且农业上的病虫害图片收集需要的成本较高，不可能拍到所有情况下的每种病虫害的图片。这就需要数据增强了，数据增强也可以看作是制造出新的数据，新的图片。

import torchvision.transforms as transforms

transform = transforms.Compose([
    transforms.RandomResizedCrop(224, scale=(0.6, 1.0)),  #随机裁剪到 224x224，保证模型输入一致
    transforms.RandomHorizontalFlip(p=0.5),  #50% 概率水平翻转，模拟不同角度的照片
    transforms.RandomRotation(30),  #旋转 -30° 到 +30° 之间，增强角度变化
    transforms.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.2),  #颜色抖动，模拟光照变化
    transforms.RandomPerspective(distortion_scale=0.5, p=0.5),  #透视变换，增加数据多样性
    transforms.RandomErasing(p=0.3, scale=(0.02, 0.2)),  #30% 概率随机遮挡图像部分区域，模拟遮挡
    transforms.ToTensor(),  #转换为 PyTorch Tensor
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  #归一化，提高训练稳定性
])

这是我的数据增强代码示例，从理论上可以让模型学会再不同光照下，各种角度，各种尺寸，遮挡的情况下也能识别。

我选择的是EfficientNet，如果用ResNet-50，需要的参数更多，计算量较多。EfficientNet需要的参数少，也相对拥有更高的准确率，在较少的数据的情况下选择这一个预训练模型更好。我一开始用的ResrNet-50的准确率在50%左右就上不去了。

import torch
import torchvision.models as models
import torch.nn as nn



#加载 EfficientNet-B3 预训练模型
model = models.efficientnet_b3(weights="IMAGENET1K_V1")  #预训练权重（在 ImageNet 训练过）

#替换分类层
num_classes = 12  #我们的病害分类数
model.classifier[1] = nn.Linear(model.classifier[1].in_features, num_classes)  #替换全连接层

其他版本的EfficientNet可以根据你数据量来选用。

import torch.optim as optim
import torch.cuda.amp as amp  #使用混合精度训练（AMP）

#定义损失函数和优化器
criterion = nn.CrossEntropyLoss()  #交叉熵损失
optimizer = optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-3)  #AdamW 优化器
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=3)  #学习率调度

scaler = amp.GradScaler()  #使用 AMP 训练

#训练循环
for epoch in range(50):
    model.train()
    running_loss = 0.0
    for inputs, labels in train_loader:
        inputs, labels = inputs.to(device), labels.to(device)

        optimizer.zero_grad()
        with amp.autocast():  #混合精度训练，提升 GPU 计算效率
            outputs = model(inputs)
            loss = criterion(outputs, labels)

        scaler.scale(loss).backward()  #反向传播
        scaler.step(optimizer)  #更新权重
        scaler.update()  #更新 `GradScaler`
        
        running_loss += loss.item()
    
    #调整学习率
    scheduler.step(running_loss)
    print(f"Epoch {epoch+1}: Loss = {running_loss:.4f}")

这里添加了自动混合精度训练，可以在不影响模型精确度的前提下加速训练，也能减少显存占用。

import matplotlib.pyplot as plt

plt.figure(figsize=(10,5))
plt.plot(history['loss'], label="Loss")  #画出 Loss 曲线
plt.plot(history['accuracy'], label="Accuracy")  #画出 Accuracy 曲线
plt.xlabel("Epoch")
plt.ylabel("Value")
plt.legend()
plt.title("Training Loss & Accuracy")
plt.savefig("training_curve.png")  #保存训练曲线
plt.show()