基于TorchGeo的遥感影像深度学习实战：从Sentinel-2到作物分类

1. 项目概述与核心价值如果你正在处理卫星影像、无人机航拍图或者任何带有地理坐标的栅格数据并且想用深度学习模型从中挖掘信息那么你很可能已经体会过那种“水土不服”的阵痛。常规的CV库如torchvision是为处理标准图片设计的它们不理解地理坐标系、投影变换也不关心如何从一张覆盖数百平方公里、包含多个波段的GeoTIFF文件中精准地切出你感兴趣的那一小块区域并与另一份矢量标注数据对齐。这正是TorchGeo要解决的核心痛点。它不是一个全新的框架而是PyTorch生态中专为地理空间数据定制的“领域库”其设计哲学与torchvision、torchaudio一脉相承让你能用熟悉的PyTorch语法来处理复杂的地理空间问题。以农业领域的作物类型制图为例这是一个典型的GeoML任务我们需要将Sentinel-2卫星拍摄的多光谱影像栅格数据与EuroCrops数据集提供的农田地块边界及作物类型标签矢量数据结合起来训练一个模型使其能够对影像中的每一个像素进行分类识别出超过200种不同的作物。这个过程远不止“读图-训练”那么简单它涉及到数据源的时空对齐、坐标系的统一转换、从大范围影像中高效采样训练小块patch、以及处理地理数据特有的I/O挑战。TorchGeo通过提供一套原生的数据集Dataset、采样器Sampler和预训练模型将这些地理空间特有的繁琐操作封装成标准的PyTorch组件让研究人员和工程师能聚焦于模型本身而非数据处理的“脏活累活”。2. 环境搭建与数据准备2.1 安装与核心依赖TorchGeo的安装非常直接得益于其完善的包管理。推荐使用Conda来管理环境以避免潜在的依赖冲突。# 创建并激活一个新的Conda环境Python 3.9-3.11兼容性较好 conda create -n torchgeo_env python3.9 conda activate torchgeo_env # 通过pip安装TorchGeo及其核心依赖 pip install torchgeo # 或者通过Conda-Forge安装有时更稳定 conda install -c conda-forge torchgeo安装命令会同时拉取PyTorch、GDAL地理空间数据读写的基础库、Rasterio、Fiona等关键依赖。如果你的环境中已经安装了PyTorchTorchGeo会自动适配。为了完成完整的训练流程我们通常还需要一些额外的库pip install segmentation-models-pytorch # 提供丰富的语义分割模型架构如U-Net pip install kornia # 提供GPU加速的数据增强操作 pip install lightning # 如果打算使用PyTorch Lightning简化训练循环注意GDAL的安装有时是最大的障碍特别是在Windows系统上。如果通过pip install torchgeo时GDAL安装失败一个更可靠的方法是先通过Conda安装GDALconda install -c conda-forge gdal然后再用pip安装TorchGeo。Conda-forge的二进制包通常能更好地解决复杂的地理库依赖。2.2 理解数据Sentinel-2与EuroCrops我们的项目需要两类数据Sentinel-2影像栅格数据这是欧洲空间局的哥白尼计划提供的多光谱卫星影像。它包含13个光谱波段可见光、红边、近红外、短波红外等空间分辨率最高达10米。对于深度学习来说这些丰富的波段提供了远超RGB图像的光谱信息是区分不同作物类型的关键。EuroCrops数据集矢量数据这是一个包含欧盟多个国家农田地块边界及其作物类型标签的矢量数据集。数据格式通常是GeoPackage或Shapefile每个多边形要素都带有“作物类型”的属性字段。项目的核心挑战在于将这两种数据“对齐”。卫星影像是覆盖整个区域的、连续的像素网格栅格而作物标签是离散的多边形矢量。训练模型时我们需要的是成对的“影像小块”和对应的“标签掩膜”。这意味着我们需要空间参考系统统一确保影像和矢量的坐标系一致。栅格化将矢量多边形转换为与影像分辨率相同的像素级标签图。采样从整个研究区域内高效、随机或有策略地抽取成千上万个固定大小如224x224像素的训练样本对。2.3 数据获取与组织TorchGeo内置了多种流行数据集的支持能自动处理下载和解压。对于Sentinel-2和EuroCrops我们可以直接指定路径或让库自动下载如果数据源支持。假设我们将数据存放在./data目录下结构如下./data/ ├── sentinel2/ # 存放Sentinel-2影像的目录可以是.tif文件 └── eurocrops/ # 存放EuroCrops矢量文件的目录可以是.gpkg或.shp文件如果使用TorchGeo的自动下载功能部分数据集支持代码会简洁很多。但更常见的场景是我们已经从其他渠道如Google Earth Engine、Copernicus Open Access Hub下载好了特定区域和时间的Sentinel-2影像。这时我们只需将文件路径提供给TorchGeo即可。3. 核心流程从原始数据到训练样本3.1 数据集创建与空间连接这是TorchGeo展示其威力的第一步。我们不再需要手动写代码去读取GeoTIFF、解析Shapefile、再进行复杂的重投影和裁剪。一切通过声明式API完成。import torch from torchgeo.datasets import Sentinel2, EuroCrops from torchgeo.samplers import RandomGeoSampler, GridGeoSampler # 1. 创建数据集对象 # Sentinel2数据集指向本地影像目录或文件列表 sentinel2_dataset Sentinel2(root./data/sentinel2) # EuroCrops数据集指定路径并允许自动下载如果本地没有 eurocrops_dataset EuroCrops(root./data/eurocrops, downloadTrue) # 2. 空间连接Spatial Join # 这是关键操作‘’运算符会计算两个数据集时空范围的交集 # 并确保返回的样本同时包含该位置的影像和标签。 # 所有重投影和栅格化都在这一步自动完成。 joint_dataset sentinel2_dataset eurocrops_dataset print(f联合数据集的空间范围: {joint_dataset.bounds}) print(f联合数据集的时间范围: {joint_dataset.bounds_t})joint_dataset现在是一个PyTorch Dataset。当你索引它时例如joint_dataset[0]它会返回一个字典其中image键对应Sentinel-2的13个波段的多光谱影像块C, H, Wmask键对应栅格化后的作物类型标签图H, W两者的空间范围完全一致。实操心得运算符背后是高效的空间索引查询。如果你的数据量很大例如整个国家的影像确保你的矢量数据和栅格数据都建立了空间索引如.csi或.sxi文件这能极大提升连接速度。对于Sentinel-2使用Cloud Optimized GeoTIFF格式能实现快速的窗口读取避免加载整张大图。3.2 数据集划分策略在机器学习中随机划分样本是常见做法。但在地理空间场景下相邻的像素或地块在空间上高度相关空间自相关简单的随机划分会导致数据泄露——即非常相似的样本同时出现在训练集和测试集使模型评估结果过于乐观。TorchGeo提供了几种地理空间感知的划分策略。from torchgeo.datasets import random_grid_cell_assignment # 在联合数据集上覆盖一个虚拟的10x10网格 # 每个网格单元被随机分配到训练、验证、测试集80%, 10%, 10% # 这种策略保证了空间上的分离评估结果更可靠。 train_dataset, val_dataset, test_dataset random_grid_cell_assignment( joint_dataset, splits[0.8, 0.1, 0.1], grid_size10 # 网格粒度可根据研究区域大小调整 )random_grid_cell_assignment是一种简单有效的方法。对于更严谨的评估你还可以使用random_bbox_assignment按地理边界框划分或自定义划分逻辑例如按行政区划或经纬度范围划分。3.3 采样器定义如何获取数据块数据集定义了“有什么数据”而采样器Sampler定义了“怎么取数据”。这是处理大范围地理影像的另一个关键。from torchgeo.samplers import RandomGeoSampler, GridGeoSampler # 训练采样器随机采样 # 在训练集范围内随机生成中心点抽取224x224大小的patch。 # 随机采样能增加数据的多样性有利于模型泛化。 train_sampler RandomGeoSampler( train_dataset, size224, # 采样patch的边长像素 length2000 # 每个epoch采样多少个patch ) # 验证和测试采样器网格采样 # 在评估区域上以224为步长无重叠地滑窗采样。 # 这样可以确保评估时覆盖整个区域且样本间不重叠便于后续拼接成完整预测图。 val_sampler GridGeoSampler( val_dataset, size224, stride224 # 步长等于size意味着无重叠 ) test_sampler GridGeoSampler(test_ds, size224, stride224)RandomGeoSampler适用于训练最大化数据多样性。GridGeoSampler适用于验证和测试确保对评估区域进行全面、无重叠的覆盖其输出可以无缝拼接回完整的地理范围用于生产完整的分类图。3.4 构建数据加载器将Dataset和Sampler组合成标准的PyTorch DataLoader就可以开始迭代训练了。from torch.utils.data import DataLoader batch_size 32 num_workers 4 # 根据CPU核心数设置用于并行数据加载 train_dataloader DataLoader( train_dataset, batch_sizebatch_size, samplertrain_sampler, num_workersnum_workers, collate_fntrain_dataset.collate_fn # TorchGeo数据集有自定义的collate函数 ) val_dataloader DataLoader(val_dataset, batch_sizebatch_size, samplerval_sampler, num_workersnum_workers, collate_fnval_dataset.collate_fn)注意事项collate_fn非常重要。地理空间数据的样本可能因为位于不同UTM带或角落而导致张量尺寸有微小差异例如由于重投影插值。TorchGeo提供的collate_fn会智能地处理这些情况确保一个batch内的数据尺寸统一。不要使用默认的collate函数。4. 模型构建、训练与评估4.1 模型定义与预训练权重加载我们选择U-Net with ResNet-50 backbone作为示例。U-Net在语义分割任务中表现出色其编码器-解码器结构能有效融合深层语义信息和浅层位置信息。TorchGeo的一个巨大优势是提供了大量在遥感数据上预训练的模型权重。import segmentation_models_pytorch as smp from torchgeo.models import ResNet50_Weights # 初始化U-Net模型 # in_channels13 对应Sentinel-2的13个波段 # classes200 对应EuroCrops中的作物类别数 model smp.Unet( encoder_nameresnet50, encoder_weightsNone, # 我们先不用ImageNet权重 in_channels13, classes200, ) # 加载TorchGeo提供的在Sentinel-2全波段上自监督预训练的权重 # 这比使用ImageNetRGB预训练权重起点高得多 weights ResNet50_Weights.SENTINEL2_ALL_DINO model.encoder.load_state_dict(weights.get_state_dict(progressTrue)) # 将模型移至GPU device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device)为什么使用遥感预训练权重ImageNet预训练模型是在RGB自然图像上训练的其卷积核学习到的特征如边缘、纹理虽然有一定通用性但对多光谱遥感影像中特有的光谱特征不敏感。TorchGeo提供的SENTINEL2_ALL_DINO等权重是通过自监督学习在海量无标签Sentinel-2影像上训练得到的其编码器已经学会了理解遥感影像的光谱、纹理和空间模式能极大加速下游任务的收敛并提升最终精度。4.2 数据预处理与增强遥感影像的原始像素值如Sentinel-2的uint16范围很大直接输入网络会导致梯度问题。我们需要进行归一化。同时数据增强对防止过拟合、提升模型鲁棒性至关重要。import kornia.augmentation as K import torch.nn as nn # 预处理将像素值从[0, 10000]Sentinel-2的反射率缩放范围归一化到[0, 1] # 注意这个归一化参数10000是Sentinel-2 L2A产品的典型值需根据你的数据产品调整。 preprocess K.Normalize(mean0.0, std10000.0) # 数据增强仅对训练数据使用 # 这里使用了简单的空间翻转。对于遥感数据还可以考虑 # - 随机旋转90180270度 # - 色彩抖动对RGB波段谨慎用于多光谱 # - 随机裁剪缩放 # 注意增强要同时应用于image和mask且mask的变换模式需设为‘nearest’以保持标签值。 augment nn.Sequential( K.RandomHorizontalFlip(p0.5), K.RandomVerticalFlip(p0.5), # K.RandomRotation(degrees90, p0.5), # 可选 )使用Kornia进行GPU上的数据增强比在CPU上使用albumentations或torchvision.transforms更快因为它能利用Tensor操作减少CPU-GPU之间的数据传输。4.3 训练循环实现下面是一个标准的PyTorch训练循环集成了上述所有组件。import torch.optim as optim from tqdm import tqdm # 用于显示进度条 # 定义损失函数和优化器 criterion nn.CrossEntropyLoss(ignore_index-1) # 忽略可能存在的无效像素 optimizer optim.AdamW(model.parameters(), lr1e-4, weight_decay1e-4) scheduler optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max100) # 余弦退火学习率调度 def train_one_epoch(epoch, dataloader): model.train() total_loss 0.0 pbar tqdm(dataloader, descfEpoch {epoch1} [Train]) for batch in pbar: # 获取数据 images batch[image].to(device) # [B, 13, H, W] masks batch[mask].to(device) # [B, H, W] # 预处理和增强仅训练时 images preprocess(images) if model.training: # Kornia增强期望输入为[B, C, H, W]且在[0,1]范围 aug_input torch.cat([images, masks.unsqueeze(1).float()], dim1) aug_output augment(aug_input) images aug_output[:, :13, :, :] masks aug_output[:, 13, :, :].long() # 前向传播 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() total_loss loss.item() pbar.set_postfix({loss: f{loss.item():.4f}}) avg_loss total_loss / len(dataloader) scheduler.step() return avg_loss def evaluate(dataloader, splitVal): model.eval() total_correct 0 total_pixels 0 with torch.no_grad(): pbar tqdm(dataloader, descf[{split}]) for batch in pbar: images batch[image].to(device) masks batch[mask].to(device) images preprocess(images) # 评估时只做预处理不做增强 outputs model(images) preds outputs.argmax(dim1) # [B, H, W] # 计算准确率忽略无效标签 valid_mask (masks ! -1) total_correct ((preds masks) valid_mask).sum().item() total_pixels valid_mask.sum().item() if split Val: # 可以在这里计算IoU等更细致的指标 pass accuracy total_correct / total_pixels if total_pixels 0 else 0 return accuracy # 主训练循环 num_epochs 100 best_val_acc 0.0 for epoch in range(num_epochs): train_loss train_one_epoch(epoch, train_dataloader) val_accuracy evaluate(val_dataloader, Val) print(fEpoch {epoch1:03d}: Train Loss {train_loss:.4f}, Val Acc {val_accuracy:.4f}) # 保存最佳模型 if val_accuracy best_val_acc: best_val_acc val_accuracy torch.save({ epoch: epoch, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), val_acc: best_val_acc, }, ./best_model.pth) print(f - Best model saved with Val Acc: {best_val_acc:.4f}) # 最终在测试集上评估 print(\n--- Final Evaluation on Test Set ---) model.load_state_dict(torch.load(./best_model.pth)[model_state_dict]) test_accuracy evaluate(test_dataloader, Test) print(fTest Set Accuracy: {test_accuracy:.4f})5. 高阶抽象使用PyTorch Lightning与CLI对于标准化的训练流程重复编写训练循环既繁琐又容易出错。TorchGeo深度集成了PyTorch Lightning提供了更高级的抽象。5.1 使用PyTorch Lightning TrainerPyTorch Lightning将训练循环、验证逻辑、日志记录、 checkpoint保存等样板代码抽象化。TorchGeo更进一步为常见地理空间任务如语义分割提供了预定义的Task和DataModule。import lightning as L from torchgeo.datamodules import Sentinel2EuroCropsDataModule from torchgeo.trainers import SemanticSegmentationTask # 1. 定义数据模块 - 它封装了之前所有的数据准备步骤 datamodule Sentinel2EuroCropsDataModule( sentinel2_paths./data/sentinel2, eurocrops_paths./data/eurocrops, batch_size32, patch_size224, num_workers4, val_split_pct0.1, test_split_pct0.1 ) # 2. 定义任务模块 - 它封装了模型、损失函数、优化器和评估指标 model SemanticSegmentationTask( modelunet, backboneresnet50, weightsResNet50_Weights.SENTINEL2_ALL_DINO, in_channels13, num_classes200, learning_rate1e-4, optimizeradamw, ) # 3. 创建Trainer并启动训练 trainer L.Trainer( max_epochs100, acceleratorauto, # 自动选择GPU/CPU devicesauto, loggerTrue, # 默认使用TensorBoard callbacks[ L.pytorch.callbacks.ModelCheckpoint(monitorval_loss, modemin), L.pytorch.callbacks.EarlyStopping(monitorval_loss, patience10), ] ) # 训练和测试 trainer.fit(model, datamoduledatamodule) trainer.test(model, datamoduledatamodule, ckpt_pathbest)使用Lightning后代码量减少了70%以上而且自动获得了模型检查点、早停、学习率监控、TensorBoard可视化等专业功能。Sentinel2EuroCropsDataModule和SemanticSegmentationTask是TorchGeo提供的“配方”对于标准任务你几乎不需要写任何训练代码。5.2 使用命令行接口进行可复现实验对于追求极致可复现性的研究或生产部署TorchGeo提供了基于LightningCLI的命令行工具。你可以将所有配置写在一个YAML文件中。# config.yaml model: class_path: torchgeo.trainers.SemanticSegmentationTask init_args: model: unet backbone: resnet50 weights: 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: ./data/sentinel2 eurocrops_paths: ./data/eurocrops batch_size: 32 patch_size: 224 num_workers: 4 val_split_pct: 0.1 test_split_pct: 0.1 trainer: max_epochs: 100 accelerator: auto devices: auto然后在终端中运行# 训练模型 torchgeo fit --config config.yaml # 使用最佳检查点测试模型 torchgeo test --config config.yaml --ckpt_path./lightning_logs/version_0/checkpoints/best.ckpt这种方式将配置与代码完全分离使得超参数搜索、实验记录和团队协作变得异常简单。任何合作者只需拥有相同的配置文件和数据集就能精确复现你的实验结果。6. 生产部署与性能优化6.1 模型导出与推理训练完成后我们需要将模型部署到生产环境进行推理。通常我们会将模型导出为TorchScript或ONNX格式以获得更好的跨平台兼容性和推理性能。# 导出为TorchScript example_input torch.randn(1, 13, 224, 224).to(device) traced_script_module torch.jit.trace(model.eval(), example_input) traced_script_module.save(crop_unet_traced.pt) # 或者导出为ONNX需要安装onnx和onnxruntime import torch.onnx torch.onnx.export( model.eval(), example_input, crop_unet.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch_size}, output: {0: batch_size}}, opset_version14 )对于地理空间推理你通常不是对单个224x224的patch感兴趣而是要对一整张大图进行预测。这需要用到GridGeoSampler的思想进行滑动窗口预测并处理重叠区域的拼接如使用高斯加权平均。6.2 性能优化与最佳实践I/O优化使用Cloud Optimized GeoTIFF。COG格式支持HTTP范围请求允许TorchGeo只读取需要的图像窗口而不是加载整个多GB的文件这对云存储和大型数据集至关重要。数据管道优化确保num_workers设置合理通常为CPU核心数。使用prefetch_factorPyTorch 1.7让DataLoader预取批次减少GPU等待时间。混合精度训练使用torch.cuda.amp进行自动混合精度训练可以显著减少GPU显存占用并加快训练速度尤其对于大型模型和输入。分布式训练对于超大数据集可以使用PyTorch的DistributedDataParallel或Lightning的distributed_backend进行多GPU或多节点训练。TorchGeo的采样器是兼容分布式训练的。缓存策略对于固定的训练集可以将预处理后的patch缓存到内存或高速SSD上如使用lmdb或WebDataset格式彻底消除I/O瓶颈。TorchGeo支持与WebDataset集成。7. 常见问题与排查技巧在实际操作中你几乎一定会遇到以下几个典型问题问题1内存溢出CUDA out of memory排查首先检查batch_size。遥感影像波段多13个patch尺寸大224x224单个样本的显存占用远高于RGB图像。使用torch.cuda.empty_cache()并尝试减小batch_size。技巧使用梯度累积。即使batch_size1也可以通过累积多个小批次的梯度再更新权重来模拟大batch的效果。技巧启用混合精度训练torch.cuda.amp可有效降低显存。问题2数据加载是训练瓶颈GPU利用率低排查使用PyTorch的torch.utils.bottleneck或nvprof工具分析。如果数据加载线程num_workers经常处于忙碌状态说明I/O或预处理太慢。解决将数据转换为COG格式。使用更快的存储NVMe SSD。将数据增强移GPU上进行使用Kornia。实现数据预取和缓存。问题3模型不收敛或准确率极低排查数据问题检查标签是否正确对齐。可视化几个样本看影像和mask是否对应。检查标签值范围0-199确认损失函数没有忽略错误的值。预处理问题确认归一化参数是否正确。Sentinel-2的表观反射率通常除以10000但如果是其他数据产品如Surface Reflectance范围可能不同。类别不平衡200类作物可能存在极端不平衡。考虑使用带权重的交叉熵损失nn.CrossEntropyLoss(weightclass_weights)其中class_weights与类别频率成反比。学习率问题尝试使用学习率查找器如Lightning的lr_find找到一个合适的初始学习率。问题4预测结果图存在明显的网格状拼接痕迹原因使用GridGeoSampler进行滑动窗口预测时窗口间无重叠导致边界处模型预测不一致。解决在推理时使用有重叠的滑动窗口并对重叠区域使用加权平均如高斯权重进行融合。这不是TorchGeo的直接功能但可以基于GridGeoSampler的思路自行实现或使用专门的推理后处理库。问题5如何在自己的数据集上使用TorchGeo核心你需要创建自定义的TorchGeoDataset。这通常意味着继承RasterDataset或VectorDataset并实现__init__,__getitem__, 和__len__方法确保你的数据能返回包含image和mask或其他键的字典以及正确的crs和bounds属性。TorchGeo的官方教程和现有数据集源码如Sentinel2类是最好的参考。通过这套从数据准备、模型训练、到高级抽象和部署的完整流程TorchGeo将地理空间深度学习的工程复杂度降到了最低。它让你能像处理普通图像一样处理卫星影像同时又不失地理信息的严谨性。无论是学术研究还是工业级应用这套工具链都提供了坚实、可扩展的基础。