AID数据集入门指南：从数据预处理到模型训练的全流程解析

1次阅读

共计 3258 个字符，预计需要花费 9 分钟才能阅读完成。

AID（Aerial Image Dataset）是一个常用的遥感图像分类数据集，包含 30 个类别的 10000 张航拍图像，每张图像大小为 600×600 像素。这个数据集非常适合用于场景分类、土地覆盖识别等计算机视觉任务。对于刚接触遥感图像处理的新手来说，AID 提供了很好的入门练习素材。

在使用 AID 数据集时，初学者经常会遇到以下几个问题：

数据格式混乱：AID 数据集原始下载的文件结构可能不够直观，需要重新组织才能方便地用于训练。
预处理效率低：由于图像尺寸较大，直接加载全部数据可能导致内存不足。
类别不平衡：某些类别的样本数量较少，可能导致模型训练时出现偏差。

首先我们需要下载并解压数据集。假设我们已经将数据集下载到 data/AID 目录下，结构如下：

# 数据集目录结构示例
AID/
├── Airport
├── BareLand
├── BaseballField
...
└── Stadium

使用 PyTorch 的 ImageFolder 可以方便地加载这种结构的数据：

from torchvision import datasets, transforms

# 基本转换
transform = transforms.Compose([transforms.Resize((256, 256)),  # 统一尺寸
    transforms.ToTensor(),          # 转为张量])

dataset = datasets.ImageFolder('data/AID', transform=transform)

查看数据集中的样本对理解数据很有帮助：

import matplotlib.pyplot as plt

# 显示 9 张图片
fig, axes = plt.subplots(3, 3, figsize=(9, 9))
for i, ax in enumerate(axes.flat):
    img, label = dataset[i]
    ax.imshow(img.permute(1, 2, 0))  # 调整通道顺序
    ax.set_title(dataset.classes[label])
    ax.axis('off')
plt.show()

除了基本的 resize 和转 tensor 外，通常还需要进行归一化：

# 计算数据集的均值和标准差
# (这里需要遍历整个数据集计算，实际中可以预先计算好)
mean = [0.485, 0.456, 0.406]  # ImageNet 常用值
std = [0.229, 0.224, 0.225]

transform = transforms.Compose([transforms.Resize((256, 256)),
    transforms.ToTensor(),
    transforms.Normalize(mean, std)
])

数据增强可以显著提升模型泛化能力，但要注意合理使用：

train_transform = transforms.Compose([transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(15),
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
    transforms.ToTensor(),
    transforms.Normalize(mean, std)
])

# 验证集不需要增强
val_transform = transforms.Compose([transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean, std)
])

常见陷阱：

过度增强可能导致模型难以收敛
不当的增强方式可能改变图像语义（如垂直翻转建筑物图像）
增强后未正确归一化

下面是一个适用于 AID 数据集的基础 CNN 模型：

import torch.nn as nn

class SimpleCNN(nn.Module):
    def __init__(self, num_classes=30):
        super(SimpleCNN, self).__init__()
        self.features = nn.Sequential(nn.Conv2d(3, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )
        self.classifier = nn.Sequential(nn.Linear(128 * 28 * 28, 512),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(512, num_classes)
        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x

import torch.optim as optim

model = SimpleCNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 数据加载器
from torch.utils.data import DataLoader
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)

# 训练循环
for epoch in range(10):  # 示例用 10 个 epoch
    model.train()
    for images, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

    # 验证
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in val_loader:
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    print(f'Epoch {epoch+1}, Val Acc: {100 * correct / total:.2f}%')