用PyTorch搞定Million-AID遥感数据集：从下载到训练，一个完整的代码仓库搭建指南

从零构建PyTorch遥感项目：Million-AID数据集全流程实战指南

引言

在计算机视觉领域，遥感图像分析正经历着前所未有的发展机遇。随着卫星和无人机技术的进步，高质量遥感数据的获取变得比以往任何时候都更加容易。然而，面对海量的遥感数据，如何有效地组织、预处理并用于深度学习模型训练，仍然是许多研究者和工程师面临的挑战。

Million-AID数据集作为当前最大的公开遥感场景分类基准之一，包含了超过100万张图像，涵盖51个精细场景类别。这个数据集不仅规模庞大，还具有复杂的层级分类结构和显著的长尾分布特性，为模型训练带来了独特的挑战。

本文将带领您从零开始，构建一个完整的PyTorch项目来处理Million-AID数据集。不同于简单的代码片段展示，我们将重点关注工程化实现和可复现性，涵盖从数据集下载、目录结构设计、自定义Dataset类实现，到训练流程优化的全流程。无论您是刚开始接触遥感图像分析的研究生，还是希望将Million-AID应用于实际项目的算法工程师，本指南都将提供可直接落地的解决方案。

1. 项目初始化与环境配置

1.1 创建项目结构

一个良好的项目结构是高效开发的基础。对于遥感图像分类任务，我们推荐以下目录组织方式：

Million-AID-Project/ ├── configs/ # 配置文件 │ └── default.yaml # 默认训练配置 ├── data/ # 数据处理模块 │ ├── __init__.py │ ├── dataset.py # 自定义Dataset类 │ ├── preprocess.py # 数据预处理 │ └── transforms.py # 自定义数据增强 ├── models/ # 模型定义 │ ├── __init__.py │ ├── base_model.py # 基础模型类 │ └── custom_cnn.py # 自定义CNN架构 ├── utils/ # 工具函数 │ ├── logger.py # 日志记录 │ └── metrics.py # 评估指标 ├── scripts/ # 实用脚本 │ ├── download_data.sh # 数据集下载 │ └── visualize.py # 数据可视化 ├── train.py # 训练入口 └── requirements.txt # 依赖列表

提示：使用tree命令可以快速生成项目目录结构图，便于文档记录和团队协作。

1.2 环境依赖安装

建议使用conda创建独立的Python环境：

conda create -n million-aid python=3.8 conda activate million-aid pip install -r requirements.txt

其中requirements.txt应包含以下核心依赖：

torch==1.12.1 torchvision==0.13.1 numpy>=1.21.0 pillow>=9.0.0 tqdm>=4.64.0 tensorboard>=2.10.0 opencv-python>=4.6.0 scikit-learn>=1.1.0

1.3 数据集获取与验证

Million-AID数据集可通过官方渠道申请下载。下载完成后，建议运行校验脚本确保数据完整性：

import hashlib import os def check_file_integrity(filepath, expected_md5): """验证文件MD5校验码""" hash_md5 = hashlib.md5() with open(filepath, "rb") as f: for chunk in iter(lambda: f.read(4096), b""): hash_md5.update(chunk) return hash_md5.hexdigest() == expected_md5 # 示例：验证训练集压缩包 train_zip_path = "Million-AID/train.zip" expected_md5 = "a1b2c3d4e5f6..." # 替换为实际MD5值 if check_file_integrity(train_zip_path, expected_md5): print("文件校验通过") else: print("文件损坏，请重新下载")

2. 高效处理Million-AID数据集

2.1 理解数据集结构

Million-AID采用三级分类体系，目录结构示例如下：

Million-AID/ ├── train/ │ ├── agriculture_land/ │ │ ├── arable_land/ │ │ │ ├── dry_field/ │ │ │ │ ├── image_001.jpg │ │ │ │ └── ... │ │ │ └── irrigated_field/ │ │ └── permanent_crop/ │ └── commercial_land/ └── test/ └── ... (类似train结构)

2.2 实现自定义Dataset类

针对这种复杂结构，我们需要设计能够自动解析类别标签的Dataset：

from torch.utils.data import Dataset from PIL import Image import os import torchvision.transforms as T class MillionAIDDataset(Dataset): def __init__(self, root_dir, transform=None, mode='train'): self.root_dir = os.path.join(root_dir, mode) self.transform = transform self.classes, self.class_to_idx = self._find_classes() self.samples = self._make_dataset() def _find_classes(self): """构建类别到索引的映射""" classes = [] for l1 in os.listdir(self.root_dir): # 第一级类别 l1_path = os.path.join(self.root_dir, l1) if not os.path.isdir(l1_path): continue for l2 in os.listdir(l1_path): # 第二级类别 l2_path = os.path.join(l1_path, l2) if not os.path.isdir(l2_path): continue for l3 in os.listdir(l2_path): # 第三级类别 l3_path = os.path.join(l2_path, l3) if os.path.isdir(l3_path): classes.append(f"{l1}/{l2}/{l3}") class_to_idx = {cls_name: i for i, cls_name in enumerate(classes)} return classes, class_to_idx def _make_dataset(self): """构建(图像路径, 标签)对的列表""" samples = [] for cls_name, cls_idx in self.class_to_idx.items(): l1, l2, l3 = cls_name.split('/') cls_dir = os.path.join(self.root_dir, l1, l2, l3) for img_name in os.listdir(cls_dir): img_path = os.path.join(cls_dir, img_name) if os.path.isfile(img_path): samples.append((img_path, cls_idx)) return samples def __len__(self): return len(self.samples) def __getitem__(self, idx): img_path, label = self.samples[idx] img = Image.open(img_path).convert('RGB') if self.transform: img = self.transform(img) return img, label

2.3 数据增强策略

针对遥感图像特点，我们设计专门的增强策略：

from torchvision import transforms # 基础变换 train_transform = transforms.Compose([ transforms.Resize(256), transforms.RandomCrop(224), transforms.RandomHorizontalFlip(), transforms.RandomVerticalFlip(), transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) # 测试集变换应保持确定性 test_transform = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ])

3. 处理长尾分布问题

3.1 分析数据分布

Million-AID存在明显的类别不平衡：

3.2 采样策略对比

我们实现多种采样器来处理不平衡问题：

from torch.utils.data import WeightedRandomSampler def get_weighted_sampler(dataset): """根据类别频率计算样本权重""" class_counts = torch.zeros(len(dataset.classes)) for _, label in dataset.samples: class_counts[label] += 1 class_weights = 1. / class_counts sample_weights = torch.tensor([class_weights[label] for _, label in dataset.samples]) return WeightedRandomSampler(sample_weights, len(sample_weights)) # 替代方案：使用现成的库 from torchsampler import ImbalancedDatasetSampler sampler = ImbalancedDatasetSampler(train_dataset)

3.3 损失函数调整

除了采样策略，还可以在损失函数层面处理不平衡：

import torch.nn as nn import torch.nn.functional as F class FocalLoss(nn.Module): def __init__(self, alpha=None, gamma=2.0): super().__init__() self.gamma = gamma self.alpha = alpha # 可传入类别权重张量 def forward(self, inputs, targets): BCE_loss = F.cross_entropy(inputs, targets, reduction='none') pt = torch.exp(-BCE_loss) loss = (1-pt)**self.gamma * BCE_loss if self.alpha is not None: alpha_t = self.alpha[targets] loss = alpha_t * loss return loss.mean() # 使用示例 alpha = torch.tensor([...]) # 根据类别频率计算 criterion = FocalLoss(alpha=alpha, gamma=2.0)

4. 模型训练与评估

4.1 基础训练流程

实现模块化的训练循环：

def train_one_epoch(model, loader, optimizer, criterion, device, epoch): model.train() total_loss = 0.0 correct = 0 total = 0 pbar = tqdm(loader, desc=f"Epoch {epoch}") for inputs, labels in pbar: inputs, labels = inputs.to(device), labels.to(device) optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() total_loss += loss.item() _, predicted = outputs.max(1) total += labels.size(0) correct += predicted.eq(labels).sum().item() pbar.set_postfix({ 'loss': total_loss/(total/loader.batch_size), 'acc': 100.*correct/total }) return total_loss/len(loader), correct/total

4.2 多尺度训练技巧

针对遥感图像特点，实现多尺度训练：

class MultiScaleTransform: """在训练过程中随机选择不同尺度""" def __init__(self, scales): self.scales = scales self.transforms = [ transforms.Compose([ transforms.Resize(scale), transforms.RandomCrop(224), # 其他变换... ]) for scale in scales ] def __call__(self, img): t = random.choice(self.transforms) return t(img) # 使用示例 train_transform = MultiScaleTransform(scales=[256, 288, 320])

4.3 模型评估指标

除了准确率，遥感任务中常用的评估指标：

from sklearn.metrics import confusion_matrix, classification_report def evaluate(model, loader, device, classes): model.eval() all_preds = [] all_labels = [] with torch.no_grad(): for inputs, labels in loader: inputs = inputs.to(device) outputs = model(inputs) _, preds = torch.max(outputs, 1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.numpy()) # 计算混淆矩阵 cm = confusion_matrix(all_labels, all_preds) print("混淆矩阵:\n", cm) # 分类报告 print("\n分类报告:") print(classification_report( all_labels, all_preds, target_names=classes )) return cm

5. 工程优化技巧

5.1 加速数据加载

使用torchdata库实现并行数据加载：

from torchdata.dataloader2 import DataLoader2, MultiProcessingReadingService def get_dataloader(dataset, batch_size, num_workers=4, shuffle=True): rs = MultiProcessingReadingService( num_workers=num_workers, worker_init_fn=lambda _: torch.manual_seed(42) ) return DataLoader2( dataset, batch_size=batch_size, shuffle=shuffle, reading_service=rs )

5.2 混合精度训练

利用AMP加速训练并减少显存占用：

from torch.cuda.amp import autocast, GradScaler scaler = GradScaler() for inputs, labels in train_loader: inputs, labels = inputs.to(device), labels.to(device) optimizer.zero_grad() with autocast(): outputs = model(inputs) loss = criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()

5.3 模型部署优化

使用TorchScript导出生产环境可用的模型：

# 导出为TorchScript model.eval() example_input = torch.rand(1, 3, 224, 224).to(device) traced_script = torch.jit.trace(model, example_input) traced_script.save("million_aid_model.pt") # 量化模型减小体积 quantized_model = torch.quantization.quantize_dynamic( model, {torch.nn.Linear}, dtype=torch.qint8 )

在实际项目中，我们发现将数据预处理逻辑也包含在TorchScript中能显著提升端到端性能。可以通过自定义nn.Module包装整个预处理+推理流程来实现这一点。