1. 项目概述与核心挑战如果你正在处理卫星影像想把一片片农田自动识别成小麦、玉米还是大豆那你大概率已经踩过地理空间机器学习GeoML的坑了。这活儿听起来酷做起来全是细节动辄几个G的GeoTIFF文件怎么高效读取不同来源的影像和矢量数据坐标系对不上怎么办如何从覆盖几百平方公里的整景影像里切出适合GPU训练的小块Patch更别提还要保证训练、验证、测试集在空间上是合理分割的不能“泄漏”信息。传统的做法是写一堆脚本用GDAL、Rasterio处理数据再用OpenCV或PIL切图最后转换成NumPy数组喂给PyTorch或TensorFlow。这个过程繁琐、易错且代码难以复用。TorchGeo的出现正是为了解决这个“最后一公里”的问题。它不是一个全新的深度学习框架而是PyTorch的一个领域库Domain Library其设计哲学是“让处理地理空间数据像处理普通图像数据一样简单”。它提供了与torchvision高度一致的API将地理空间数据特有的概念如坐标参考系统CRS、边界框Bounding Box封装成PyTorch友好的Dataset和Sampler。这意味着你可以用写普通CV项目的思维和代码结构来处理带有地理坐标的卫星影像。本次实战我们将聚焦一个经典且实用的任务基于Sentinel-2多光谱影像的欧盟作物类型精细分类。数据方面我们使用Sentinel-2 L2A级地表反射率产品作为输入栅格数据以及EuroCrops数据集提供的农田地块边界及作物类型标签矢量数据。目标是训练一个模型为影像中的每个像素预测其所属的作物类别约200类。我们将使用带有ResNet-50编码器的U-Net架构这是一个在语义分割任务中久经考验的模型。整个流程将贯穿GeoML项目的核心环节从原始数据加载与对齐、空间采样策略设计、深度学习模型构建与训练到最终模型评估。我会在每一步都分享我趟过的坑和总结的经验让你不仅能复现这个项目更能理解其背后的设计逻辑从而应用到你自己的地理空间问题中。2. 环境搭建与数据准备2.1 安装TorchGeo与依赖库TorchGeo的安装非常直接因为它已经上架了PyPI和Conda Forge。我强烈建议在一个新的虚拟环境中进行以避免包冲突。# 使用pip安装 pip install torchgeo # 或者使用conda安装 conda install torchgeo -c conda-forge安装TorchGeo会同时安装其核心依赖包括PyTorch、Rasterio、Fiona、GeoPandas等。但为了完成完整的训练流程我们还需要一些额外的库pip install segmentation-models-pytorch kornia lightningsegmentation-models-pytorch (smp)提供了预构建的U-Net、DeepLabV3等分割模型 backbone 选择丰富非常方便。kornia一个基于PyTorch的计算机视觉库它的数据增强操作直接在GPU张量上进行速度远超基于PIL/OpenCV的增强并且能完美集成到PyTorch数据流中。lightning (PyTorch Lightning)用于组织训练代码让代码更简洁、模块化并内置了TensorBoard日志、模型检查点等实用功能。注意确保你的PyTorch版本与CUDA版本匹配如果使用GPU。可以先去 PyTorch官网 获取正确的安装命令再安装TorchGeo。2.2 理解数据Sentinel-2与EuroCrops我们的任务需要两类数据Sentinel-2 L2A 影像栅格数据这是欧空局哥白尼计划下的多光谱卫星数据包含13个光谱波段可见光、红边、近红外、短波红外等空间分辨率最高10米。L2A级产品是经过大气校正的地表反射率数据更适合用于土地覆盖分类。数据通常以“景”Scene为单位每景覆盖约100x100公里存储为Cloud Optimized GeoTIFF (COG)格式。COG格式是关键它允许我们随机读取大文件中的一小块区域而无需加载整个文件到内存这对处理海量遥感数据至关重要。EuroCrops 数据集矢量数据这是一个包含欧盟多个国家农田地块边界及其官方申报作物类型的数据集。数据以GeoPackage或Shapefile格式提供每个多边形要素Polygon代表一块农田并包含一个class_id属性对应具体的作物类型如1代表冬小麦2代表玉米等。核心挑战与TorchGeo的解决方案空间对齐Sentinel-2影像和EuroCrops矢量数据可能使用不同的坐标参考系统CRS。手动处理需要先用GDAL进行重投影Reprojection非常麻烦。TorchGeo的Dataset在加载数据时会自动将所有数据动态重投影到一个统一的CRS默认是第一个数据集的CRS这个操作对用户透明。数据配对我们需要找到每个EuroCrops多边形所对应的Sentinel-2影像区域即空间交集。TorchGeo的操作符交集和|操作符并集可以优雅地处理多个数据集的空间-时间索引自动找出时空重叠的部分。标签生成EuroCrops是矢量多边形而我们需要的是与影像像素一一对应的分类标签图Mask。TorchGeo会在需要时动态地将矢量数据栅格化Rasterize生成与当前影像Patch大小和空间范围完全匹配的标签Mask。2.3 数据下载与组织假设你已经从相关平台如Copernicus Open Access Hub, EuroCrops官网下载了部分数据。你需要将数据组织成TorchGeo能识别的结构。通常Sentinel-2数据按景存放EuroCrops数据按国家或区域存放。一个简单的目录结构示例如下your_data_root/ ├── sentinel2/ │ ├── T32TNS_20230501T100031_B02.jp2 │ ├── T32TNS_20230501T100031_B03.jp2 │ └── ... (其他波段文件) └── eurocrops/ ├── DE_2023.gpkg # 德国数据 └── FR_2023.gpkg # 法国数据TorchGeo的Sentinel2和EuroCrops数据集类可以接受一个包含文件路径的列表或者一个根目录路径它会根据内置规则搜索文件。对于EuroCrops设置downloadTrue可以自动下载数据如果本地没有但这需要稳定的网络连接和较大的磁盘空间。对于生产环境建议先手动下载并管理数据。3. 核心流程解析从数据加载到模型训练TorchGeo提供了三种不同抽象级别的编程范式从高控制度的纯PyTorch到便捷的命令行。我们从最灵活、最能体现其设计思想的“纯PyTorch”方式开始。3.1 纯PyTorch范式最大程度的控制这种方式直接使用TorchGeo提供的Dataset、Sampler与标准的PyTorch训练循环无缝集成。3.1.1 数据集创建与空间索引首先导入必要的模块并创建数据集对象。import kornia.augmentation as K import segmentation_models_pytorch as smp import torch from torch import nn, optim from torch.utils.data import DataLoader from torchgeo.datasets import Sentinel2, EuroCrops, random_grid_cell_assignment from torchgeo.samplers import GridGeoSampler, RandomGeoSampler from torchgeo.models import ResNet50_Weights # 1. 创建数据集 # 假设你的数据存放在本地路径 sentinel2_root ./your_data_root/sentinel2/ eurocrops_root ./your_data_root/eurocrops/ # 初始化数据集。Sentinel2会自动识别波段并堆叠。 # paths参数可以是目录路径也可以是具体文件列表。 sentinel2 Sentinel2(pathssentinel2_root) # EuroCrops数据集设置downloadFalse因为我们已手动准备 eurocrops EuroCrops(pathseurocrops_root, downloadFalse) # 2. 计算时空交集 # 这是关键一步操作符会基于空间边界框Bounding Box和可选的时间范围 # 自动找出两个数据集重叠的区域并返回一个新的联合数据集。 # 这个新数据集在调用时会同时返回对应区域的影像(image)和标签(mask)。 dataset sentinel2 eurocrops print(f”联合数据集共有 {len(dataset)} 个时空查询索引。”)这里发生了什么dataset对象现在是一个IntersectionDataset。它内部维护了一个空间索引R-tree记录了所有sentinel2和eurocrops数据在空间和时间上的重叠部分。当你通过Sampler询问某个位置的数据时它会分别从两个原始数据集中加载对应区域的影像和矢量数据并自动完成重投影和栅格化确保返回的image和mask在空间上严格对齐。3.1.2 地理空间采样策略在普通图像数据集中我们通常随机打乱所有图片。但在地理空间数据中相邻的像素或Patch在空间上是自相关的即地理学第一定律万物皆相关相近的事物关联更紧密。如果简单随机划分训练/验证/测试集会导致数据泄漏模型在训练时“见过”验证集附近的信息从而高估模型性能。TorchGeo提供了多种地理空间感知的数据划分策略。这里我们使用random_grid_cell_assignment# 3. 地理空间数据划分 # 在数据集上覆盖一个10x10的虚拟网格。 # 每个网格单元被随机分配到训练、验证或测试集比例为80:10:10。 # 这种策略确保了不同集合的样本在空间上是分离的避免了泄漏。 train_ds, val_ds, test_ds random_grid_cell_assignment( dataset, splits[0.8, 0.1, 0.1], grid_size10 )接下来为每个划分创建采样器Sampler。采样器的核心作用是给定一个地理空间数据集决定从哪些位置以边界框表示采集固定大小的Patch。# 4. 创建地理空间采样器 patch_size 224 # 像素单位对应Sentinel-2 10米分辨率即2.24km x 2.24km的区域 # 训练时使用随机采样增加数据多样性 train_sampler RandomGeoSampler(train_ds, sizepatch_size, length10000) # length决定每个epoch采样的Patch数量 # 验证和测试时使用网格采样确保无重叠、无遗漏地覆盖整个区域便于做整体评估 val_sampler GridGeoSampler(val_ds, sizepatch_size, stridepatch_size) test_sampler GridGeoSampler(test_ds, sizepatch_size, stridepatch_size)RandomGeoSampler在数据集的空间范围内随机生成边界框。length参数很重要它定义了每个epoch要采样多少个Patch。你可以根据数据集大小和批次大小来设定。GridGeoSampler以固定的步长stride在数据集空间范围内滑动生成一个覆盖整个区域的、无重叠的Patch网格。这通常用于模型推理或生成完整的预测图。3.1.3 构建数据加载器与数据预处理采样器准备好后就可以和PyTorch标准的DataLoader结合了。# 5. 创建PyTorch DataLoader batch_size 16 # 根据GPU内存调整 train_loader DataLoader(train_ds, batch_sizebatch_size, samplertrain_sampler, num_workers4) val_loader DataLoader(val_ds, batch_sizebatch_size, samplerval_sampler, num_workers4) test_loader DataLoader(test_ds, batch_sizebatch_size, samplertest_sampler, num_workers4)num_workers用于设置多进程数据加载可以显著加速IO密集的数据读取过程。在Linux系统上效果显著在Windows上可能需要设置为0。接下来是预处理和数据增强。Sentinel-2 L2A数据的反射率值范围通常是0-10000或0-1取决于处理级别。我们需要将其归一化。# 6. 定义预处理和数据增强 # Kornia的变换操作接收和返回的都是PyTorch Tensor且支持GPU加速。 preprocess K.Normalize(mean0.0, std10000.0) # 将[0, 10000]范围归一化到约[0,1] # 数据增强简单的水平和垂直翻转。对于遥感影像旋转需要谨慎可能改变地理方位。 augment K.ImageSequential( K.RandomHorizontalFlip(p0.5), K.RandomVerticalFlip(p0.5), # K.RandomRotation(degrees90, p0.5), # 谨慎使用旋转 data_keys[“input”] # 指定对’image‘进行增强’mask‘需要相应的几何变换稍复杂 ) # 注意上述增强只针对’image‘。对于语义分割任务’mask‘需要同步进行完全相同的空间变换。 # TorchGeo目前不直接处理这个同步一种常见做法是自定义Dataset或使用Albumentations库。 # 为简化本例暂不对mask做增强或使用smp或Albumentations中支持同步变换的增强管道。实操心得对于多光谱遥感影像波段顺序至关重要。Sentinel-2默认的波段顺序可能与某些预训练模型期望的顺序不同。TorchGeo的Sentinel2数据集默认会按照波长顺序B02, B03, B04, B08, ...返回13个波段。务必确认你的模型输入通道数与数据一致。此外数据增强时切勿对影像进行颜色抖动或亮度对比度调整这会破坏地表反射率的物理意义。空间变换翻转、裁剪是安全的。3.1.4 模型构建与权重加载我们使用segmentation-models-pytorch来快速构建一个U-Net。# 7. 构建模型 # 输入通道为13Sentinel-2所有波段输出类别为EuroCrops中的作物种类数约200 model smp.Unet( encoder_nameresnet50, encoder_weightsNone, # 我们先不用ImageNet权重 in_channels13, classes200, ) # 8. 加载针对遥感数据预训练的权重关键步骤 # TorchGeo提供了在Sentinel-2数据上通过自监督学习如DINO预训练的ResNet权重。 # 这比使用ImageNet预训练权重基于自然图像的起点要好得多。 weights ResNet50_Weights.SENTINEL2_ALL_DINO # 只加载编码器encoder部分的权重 model.encoder.load_state_dict(weights.get_state_dict(progressTrue)) # 冻结编码器的前几层进行微调可选 for param in model.encoder.parameters(): param.requires_grad False # 可以解冻最后几层 for param in list(model.encoder.children())[-2:]: for p in param.parameters(): p.requires_grad True device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device)使用在遥感数据上预训练的骨干网络是提升GeoML模型性能的关键技巧。自然图像ImageNet和卫星影像的纹理、特征分布存在差异使用领域预训练权重能加速收敛并获得更好效果。3.1.5 训练循环与评估剩下的就是标准的PyTorch训练流程。# 9. 定义损失函数和优化器 criterion nn.CrossEntropyLoss(ignore_index-1) # 忽略标签为-1的像素可能为无数据区域 optimizer optim.AdamW(model.parameters(), lr1e-4, weight_decay1e-4) scheduler optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max50) def train_one_epoch(loader, model, criterion, optimizer, device): model.train() running_loss 0.0 for batch in loader: images batch[image].to(device) # [B, 13, H, W] masks batch[mask].to(device).long() # [B, H, W] # 预处理 images preprocess(images) # 数据增强仅对图像 images augment(images) # 前向传播 optimizer.zero_grad() outputs model(images) # [B, 200, H, W] loss criterion(outputs, masks) # 反向传播 loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) # 梯度裁剪 optimizer.step() running_loss loss.item() * images.size(0) epoch_loss running_loss / len(loader.dataset) return epoch_loss def evaluate(loader, model, device): model.eval() total_correct 0 total_pixels 0 with torch.no_grad(): for batch in loader: images batch[image].to(device) masks batch[mask].to(device).long() images preprocess(images) outputs model(images) preds outputs.argmax(dim1) # [B, H, W] valid_mask masks ! -1 # 忽略无数据区域 total_correct (preds[valid_mask] masks[valid_mask]).sum().item() total_pixels valid_mask.sum().item() accuracy total_correct / total_pixels if total_pixels 0 else 0 return accuracy # 10. 训练循环 num_epochs 50 for epoch in range(num_epochs): train_loss train_one_epoch(train_loader, model, criterion, optimizer, device) val_acc evaluate(val_loader, model, device) scheduler.step() print(f”Epoch {epoch1}/{num_epochs}, Loss: {train_loss:.4f}, Val Acc: {val_acc:.4f}”) # 这里可以添加模型保存、早停等逻辑3.2 PyTorch Lightning范式简化训练代码如果你觉得上面的训练循环样板代码太多TorchGeo与PyTorch Lightning的集成可以大幅简化代码。PyTorch Lightning将训练逻辑训练/验证/测试步骤、优化器调度、日志记录等抽象成LightningModule并将数据加载抽象成LightningDataModule。TorchGeo为许多常见数据集和任务提供了现成的DataModule和Task即LightningModule。from lightning.pytorch import Trainer from torchgeo.datamodules import Sentinel2EuroCropsDataModule from torchgeo.trainers import SemanticSegmentationTask # 1. 定义数据模块 datamodule Sentinel2EuroCropsDataModule( sentinel2_paths./your_data_root/sentinel2/, eurocrops_paths./your_data_root/eurocrops/, batch_size16, patch_size224, num_workers4, val_split_pct0.1, # 验证集比例 test_split_pct0.1, # 测试集比例 ) # 2. 定义任务模型训练逻辑 model SemanticSegmentationTask( modelunet, backboneresnet50, weightssentinel2_all_dino, # 直接使用字符串指定权重 in_channels13, num_classes200, learning_rate1e-4, optimizeradamw, ) # 3. 创建训练器并训练 trainer Trainer( max_epochs50, acceleratorgpu if torch.cuda.is_available() else cpu, devices1, loggerTrue, # 默认使用TensorBoard enable_checkpointingTrue, ) trainer.fit(modelmodel, datamoduledatamodule) trainer.test(modelmodel, datamoduledatamodule)使用PyTorch Lightning后代码量减少了约70%。Trainer自动处理了训练循环、验证、日志记录、模型保存Checkpointing、甚至分布式训练。Sentinel2EuroCropsDataModule内部封装了之前所有的数据加载、交集计算、划分和采样逻辑。SemanticSegmentationTask则封装了模型构建、损失函数交叉熵、以及评估指标如准确率、IoU。3.3 命令行范式追求极致复现性对于研究或需要频繁实验的场景将配置与代码分离是最佳实践。TorchGeo提供了基于LightningCLI的命令行接口允许你通过YAML或JSON文件来定义整个实验。创建一个config.yaml文件# config.yaml model: class_path: torchgeo.trainers.SemanticSegmentationTask init_args: model: unet backbone: resnet50 weights: sentinel2_all_dino in_channels: 13 num_classes: 200 learning_rate: 1e-4 optimizer: adamw data: class_path: torchgeo.datamodules.Sentinel2EuroCropsDataModule init_args: sentinel2_paths: ./your_data_root/sentinel2/ eurocrops_paths: ./your_data_root/eurocrops/ batch_size: 16 patch_size: 224 num_workers: 4 val_split_pct: 0.1 test_split_pct: 0.1 trainer: max_epochs: 50 accelerator: auto devices: 1 logger: true enable_checkpointing: true然后在命令行中运行# 训练模型 torchgeo fit --config config.yaml # 使用训练好的检查点进行测试 torchgeo test --config config.yaml --ckpt_path./lightning_logs/version_0/checkpoints/last.ckpt这种方式将实验配置完全固化在文件中任何人拿到配置文件和相同的数据都能一键复现你的训练过程极大提升了研究的可复现性。4. 实战中的关键问题与解决方案4.1 数据不均衡与类别权重作物分类任务中类别极度不均衡是常态例如大面积的小麦和零星的特种作物。直接使用交叉熵损失会导致模型偏向主导类别。解决方案样本加权在RandomGeoSampler中可以根据每个空间位置或网格单元的类别分布赋予不同的采样概率让模型更多地看到稀有类别的区域。损失函数加权为交叉熵损失函数中的每个类别计算权重。权重通常与类别频率成反比。# 示例计算类别权重需要在数据加载后统计 from collections import Counter import numpy as np def compute_class_weights(loader, num_classes): pixel_counts Counter() for batch in loader: masks batch[mask].numpy().flatten() # 忽略无效值如-1 masks masks[masks 0] pixel_counts.update(masks) total_pixels sum(pixel_counts.values()) class_weights np.zeros(num_classes) for cls, count in pixel_counts.items(): if count 0: # 使用中位数频率平衡法 freq count / total_pixels class_weights[cls] 1 / (np.log(1.2 freq)) # 一种平滑的逆频率加权 # 归一化权重 class_weights torch.from_numpy(class_weights / class_weights.sum()).float().to(device) return class_weights # 使用加权损失 class_weights compute_class_weights(train_loader, num_classes200) criterion nn.CrossEntropyLoss(weightclass_weights, ignore_index-1)4.2 处理大范围区域与内存限制当研究区域很大时无法一次性将整个区域的标签Mask栅格化到内存中。TorchGeo的动态栅格化机制完美解决了这个问题。只有在Sampler请求某个特定Patch时它才会在该Patch的边界框范围内将矢量数据栅格化内存消耗只与Patch大小有关与总体区域大小无关。4.3 多时相数据融合作物生长具有强烈的季节性。仅用单一时相的影像很难区分某些作物如春小麦和冬小麦。TorchGeo的IntersectionDataset同样支持时间维度的交集计算。你可以创建包含多个时相Sentinel-2数据的数据集然后与EuroCrops取交集。返回的image将是一个[B, T, C, H, W]的张量T为时相数。模型需要相应调整以处理时序数据例如使用3D卷积或Transformer。4.4 模型选择与调优Backbone选择除了ResNet-50可以尝试更高效的架构如EfficientNet、ConvNeXt或在遥感数据上预训练的模型如ResNet50_Weights.SENTINEL2_ALL_MOCO。解码器选择U-Net是经典选择。对于细节要求更高的任务可以尝试DeepLabV3具有空洞空间金字塔池化或FPN特征金字塔网络。输入波段选择Sentinel-2有13个波段但并非所有都对作物分类有用。可以尝试只使用RGB和近红外波段B2, B3, B4, B8或者加入红边波段B5, B6, B7以观察效果。使用波段索引功能可以轻松选择。sentinel2 Sentinel2(paths..., bands[“B02”, “B03”, “B04”, “B08”]) # 只加载4个波段4.5 评估指标与可视化除了整体准确率在语义分割中更常用的指标是交并比IoU和平均IoUmIoU。可以使用torchmetrics库方便地计算。from torchmetrics.classification import MulticlassJaccardIndex num_classes 200 iou_metric MulticlassJaccardIndex(num_classesnum_classes, ignore_index-1).to(device) def evaluate_iou(loader, model, device): model.eval() iou_metric.reset() with torch.no_grad(): for batch in loader: images batch[image].to(device) masks batch[mask].to(device).long() images preprocess(images) outputs model(images) preds outputs.argmax(dim1) iou_metric.update(preds, masks) miou iou_metric.compute() return miou可视化对于调试至关重要。可以定期保存模型对验证集Patch的预测结果并与真值对比直观查看模型在哪些作物、哪些区域上表现不佳。5. 项目总结与进阶方向通过这个完整的流程我们利用TorchGeo实现了从原始地理空间数据到深度学习模型训练的端到端管道。其核心价值在于将复杂的地理空间数据处理抽象成标准的PyTorch数据流让研究者能更专注于模型设计和算法创新而非数据处理的“脏活累活”。回顾几个关键点数据即代码TorchGeo的Dataset对象封装了数据读取、重投影、栅格化等所有预处理保证了数据的一致性。空间感知的采样与划分这是GeoML区别于普通CV的核心避免了数据泄漏评估结果更可信。领域预训练权重使用在遥感数据上预训练的模型作为起点是提升性能的捷径。灵活的抽象层级从底层PyTorch API到高层命令行工具满足了从研究原型到生产部署的不同需求。未来的进阶方向扩展到更多传感器尝试融合Sentinel-1 SAR数据对云层不敏感对作物结构敏感或Landsat、MODIS等数据。时序模型引入LSTM、Transformer或3D CNN来处理时间序列数据捕捉作物物候变化。弱监督与半监督学习标注地理空间数据成本高昂。可以利用TorchGeo方便地集成自监督学习如MoCo, DINO预训练或使用点标注、边界框等弱标签进行训练。模型部署与推理将训练好的模型应用于新的、大范围的卫星影像生成作物分类图。需要考虑使用GridGeoSampler进行滑动窗口推理并处理相邻Patch之间的接缝问题如使用重叠窗口和加权融合。TorchGeo生态仍在快速发展其模块化设计使得它易于与新的模型架构、学习范式集成。对于任何希望将深度学习应用于地理空间数据的研究者或工程师来说掌握TorchGeo无疑将极大地提升工作效率和项目质量。
