30万张图像、1081类植物的细分类数据集

关注

原创 2026-08-13 17:53:54 229 阅读 AI 写同款· GEO 优化 ›

深度学习11 — ResNet-50 算法实战与解析 ResNet-50 算法实战与解析阅读详情

在这里插入图片描述
对于一个包含30万张图像、1081类植物的细分类数据集，您需要设置一个高效的图像分类系统。这个数据集具有高度内在歧义和长尾分布的特点，因此在设计模型时需要特别注意这些挑战。以下是如何准备和训练这样一个大规模分类数据集的步骤。、在这里插入图片描述

30万图像1081类植物细分类数据集，注意是分类数据，没有检测框信息共33GB，该数据集具有高度内在歧义和长尾分布，可用于细分类识别任务，含训练代码在这里插入图片描述

1. 环境准备

确保安装了必要的库和工具：

pip install torch torchvision
pip install numpy
pip install matplotlib
pip install opencv-python
pip install pyyaml
pip install albumentations
pip install scikit-learn

2. 数据集准备

假设您的数据集目录结构如下：

plant_classification_dataset/
├── train/
│   ├── class_0/
│   ├── class_1/
│   └── ...
├── val/
│   ├── class_0/
│   ├── class_1/
│   └── ...
└── plant_data.yaml

每个子文件夹代表一个类别，其中包含了该类别的所有图像。

3. 创建数据集配置文件

创建一个 plant_data.yaml 文件，内容如下：

train: ../plant_classification_dataset/train
val: ../plant_classification_dataset/val

nc: 1081  # 总共有1081个类别
names: ['class_0', 'class_1', ..., 'class_1080']  # 填充所有类别名称

4. 准备数据加载器

使用PyTorch的数据加载器来读取图像和标签，并进行数据增强。

import os
from torch.utils.data import Dataset, DataLoader
import cv2
import numpy as np
from sklearn.model_selection import train_test_split
from torchvision import transforms
from PIL import Image

class PlantClassificationDataset(Dataset):
    def __init__(self, data_dir, transform=None):
        self.data_dir = data_dir
        self.transform = transform
        self.classes = sorted(os.listdir(data_dir))
        self.class_to_idx = {cls_name: i for i, cls_name in enumerate(self.classes)}
        self.images = []
        self.labels = []

        for cls_name in self.classes:
            cls_dir = os.path.join(data_dir, cls_name)
            for img_name in os.listdir(cls_dir):
                img_path = os.path.join(cls_dir, img_name)
                self.images.append(img_path)
                self.labels.append(self.class_to_idx[cls_name])

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        img_path = self.images[idx]
        image = Image.open(img_path).convert('RGB')
        label = self.labels[idx]

        if self.transform:
            image = self.transform(image)

        return image, label

# 数据增强
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(15),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 创建数据加载器
train_dataset = PlantClassificationDataset(train_data_dir, transform=transform)
val_dataset = PlantClassificationDataset(val_data_dir, transform=transform)

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4)

5. 定义模型

您可以选择多种预训练模型作为基础，例如ResNet、EfficientNet等。这里我们以ResNet50为例。

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

class PlantClassifier(nn.Module):
    def __init__(self, num_classes=1081):
        super(PlantClassifier, self).__init__()
        self.base_model = models.resnet50(pretrained=True)
        self.base_model.fc = nn.Linear(self.base_model.fc.in_features, num_classes)

    def forward(self, x):
        return self.base_model(x)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = PlantClassifier(num_classes=1081).to(device)

6. 训练模型

编写训练循环来训练模型。

import torch.optim as optim
from torch.utils.tensorboard import SummaryWriter
from datetime import datetime

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-4)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)

writer = SummaryWriter(f"runs/{datetime.now().strftime('%Y%m%d-%H%M%S')}")

num_epochs = 50
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for i, (images, labels) in enumerate(train_loader):
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()

        if (i + 1) % 10 == 0:
            print(f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{len(train_loader)}], Loss: {loss.item():.4f}')
            writer.add_scalar('Training Loss', loss.item(), epoch * len(train_loader) + i)

    # 验证阶段
    model.eval()
    with torch.no_grad():
        correct = 0
        total = 0
        for images, labels in val_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
        val_accuracy = 100 * correct / total
        print(f'Validation Accuracy: {val_accuracy:.2f}%')
        writer.add_scalar('Validation Accuracy', val_accuracy, epoch)

    scheduler.step()

    # 保存模型
    torch.save(model.state_dict(), f'models/plant_classifier_epoch_{epoch+1}.pth')

writer.close()

7. 评估模型

训练完成后，您可以使用以下代码来评估模型在验证集上的性能。

def evaluate_model(model, dataloader, device):
    model.eval()
    with torch.no_grad():
        correct = 0
        total = 0
        for images, labels in dataloader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
        accuracy = 100 * correct / total
        print(f'Accuracy: {accuracy:.2f}%')

evaluate_model(model, val_loader, device)