基于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 python=3.9 conda activate torchgeo_env # 通过pip安装TorchGeo及其核心依赖 pip install torchgeo # 或者通过Conda-Forge安装（有时更稳定） conda install -c conda-forge torchgeo

安装命令会同时拉取PyTorch、GDAL（地理空间数据读写的基础库）、Rasterio、Fiona等关键依赖。如果你的环境中已经安装了PyTorch，TorchGeo会自动适配。为了完成完整的训练流程，我们通常还需要一些额外的库：

pip install segmentation-models-pytorch # 提供丰富的语义分割模型架构，如U-Net pip install kornia # 提供GPU加速的数据增强操作 pip install lightning # 如果打算使用PyTorch Lightning简化训练循环

注意：GDAL的安装有时是最大的障碍，特别是在Windows系统上。如果通过pip install torchgeo时GDAL安装失败，一个更可靠的方法是先通过Conda安装GDAL：conda 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", download=True) # 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, W），'mask'键对应栅格化后的作物类型标签图（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_size=10 # 网格粒度，可根据研究区域大小调整 )

random_grid_cell_assignment是一种简单有效的方法。对于更严谨的评估，你还可以使用random_bbox_assignment（按地理边界框划分）或自定义划分逻辑，例如按行政区划或经纬度范围划分。

3.3 采样器：定义如何获取数据块

数据集定义了“有什么数据”，而采样器（Sampler）定义了“怎么取数据”。这是处理大范围地理影像的另一个关键。

from torchgeo.samplers import RandomGeoSampler, GridGeoSampler # 训练采样器：随机采样 # 在训练集范围内随机生成中心点，抽取224x224大小的patch。 # 随机采样能增加数据的多样性，有利于模型泛化。 train_sampler = RandomGeoSampler( train_dataset, size=224, # 采样patch的边长（像素） length=2000 # 每个epoch采样多少个patch ) # 验证和测试采样器：网格采样 # 在评估区域上以224为步长，无重叠地滑窗采样。 # 这样可以确保评估时覆盖整个区域，且样本间不重叠，便于后续拼接成完整预测图。 val_sampler = GridGeoSampler( val_dataset, size=224, stride=224 # 步长等于size，意味着无重叠 ) test_sampler = GridGeoSampler(test_ds, size=224, stride=224)

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_size=batch_size, sampler=train_sampler, num_workers=num_workers, collate_fn=train_dataset.collate_fn # TorchGeo数据集有自定义的collate函数 ) val_dataloader = DataLoader(val_dataset, batch_size=batch_size, sampler=val_sampler, num_workers=num_workers, collate_fn=val_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_channels=13 对应Sentinel-2的13个波段 # classes=200 对应EuroCrops中的作物类别数 model = smp.Unet( encoder_name='resnet50', encoder_weights=None, # 我们先不用ImageNet权重 in_channels=13, classes=200, ) # 加载TorchGeo提供的在Sentinel-2全波段上自监督预训练的权重 # 这比使用ImageNet（RGB）预训练权重起点高得多！ weights = ResNet50_Weights.SENTINEL2_ALL_DINO model.encoder.load_state_dict(weights.get_state_dict(progress=True)) # 将模型移至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(mean=0.0, std=10000.0) # 数据增强：仅对训练数据使用 # 这里使用了简单的空间翻转。对于遥感数据，还可以考虑： # - 随机旋转（90，180，270度） # - 色彩抖动（对RGB波段，谨慎用于多光谱） # - 随机裁剪缩放 # 注意：增强要同时应用于image和mask，且mask的变换模式需设为‘nearest’以保持标签值。 augment = nn.Sequential( K.RandomHorizontalFlip(p=0.5), K.RandomVerticalFlip(p=0.5), # K.RandomRotation(degrees=90, p=0.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(), lr=1e-4, weight_decay=1e-4) scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100) # 余弦退火学习率调度 def train_one_epoch(epoch, dataloader): model.train() total_loss = 0.0 pbar = tqdm(dataloader, desc=f'Epoch {epoch+1} [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()], dim=1) 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_norm=1.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, split='Val'): model.eval() total_correct = 0 total_pixels = 0 with torch.no_grad(): pbar = tqdm(dataloader, desc=f'[{split}]') for batch in pbar: images = batch['image'].to(device) masks = batch['mask'].to(device) images = preprocess(images) # 评估时只做预处理，不做增强 outputs = model(images) preds = outputs.argmax(dim=1) # [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(f'Epoch {epoch+1: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(f'Test Set Accuracy: {test_accuracy:.4f}')

5. 高阶抽象：使用PyTorch Lightning与CLI

对于标准化的训练流程，重复编写训练循环既繁琐又容易出错。TorchGeo深度集成了PyTorch Lightning，提供了更高级的抽象。

5.1 使用PyTorch Lightning Trainer

PyTorch 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_size=32, patch_size=224, num_workers=4, val_split_pct=0.1, test_split_pct=0.1 ) # 2. 定义任务模块 - 它封装了模型、损失函数、优化器和评估指标 model = SemanticSegmentationTask( model='unet', backbone='resnet50', weights=ResNet50_Weights.SENTINEL2_ALL_DINO, in_channels=13, num_classes=200, learning_rate=1e-4, optimizer='adamw', ) # 3. 创建Trainer并启动训练 trainer = L.Trainer( max_epochs=100, accelerator='auto', # 自动选择GPU/CPU devices='auto', logger=True, # 默认使用TensorBoard callbacks=[ L.pytorch.callbacks.ModelCheckpoint(monitor='val_loss', mode='min'), L.pytorch.callbacks.EarlyStopping(monitor='val_loss', patience=10), ] ) # 训练和测试 trainer.fit(model, datamodule=datamodule) trainer.test(model, datamodule=datamodule, ckpt_path='best')

使用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_version=14 )

对于地理空间推理，你通常不是对单个224x224的patch感兴趣，而是要对一整张大图进行预测。这需要用到GridGeoSampler的思想，进行滑动窗口预测，并处理重叠区域的拼接（如使用高斯加权平均）。

6.2 性能优化与最佳实践

I/O优化：使用Cloud Optimized GeoTIFF。COG格式支持HTTP范围请求，允许TorchGeo只读取需要的图像窗口，而不是加载整个多GB的文件，这对云存储和大型数据集至关重要。
数据管道优化：确保num_workers设置合理（通常为CPU核心数）。使用prefetch_factor（PyTorch 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_size=1，也可以通过累积多个小批次的梯度再更新权重，来模拟大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(weight=class_weights)），其中class_weights与类别频率成反比。
- 学习率问题：尝试使用学习率查找器（如Lightning的lr_find）找到一个合适的初始学习率。

问题4：预测结果图存在明显的网格状拼接痕迹

原因：使用GridGeoSampler进行滑动窗口预测时，窗口间无重叠，导致边界处模型预测不一致。
解决：在推理时使用有重叠的滑动窗口，并对重叠区域使用加权平均（如高斯权重）进行融合。这不是TorchGeo的直接功能，但可以基于GridGeoSampler的思路自行实现，或使用专门的推理后处理库。

问题5：如何在自己的数据集上使用TorchGeo？

核心：你需要创建自定义的TorchGeoDataset。这通常意味着继承RasterDataset或VectorDataset，并实现__init__,__getitem__, 和__len__方法，确保你的数据能返回包含image和mask（或其他键）的字典，以及正确的crs和bounds属性。TorchGeo的官方教程和现有数据集源码（如Sentinel2类）是最好的参考。