TorchGeo 是一个基于 PyTorch 的领域专用库，类似于 torchvision，为地理空间数据提供了专门的数据集、采样器、变换和预训练模型。

该库的目标是让以下两件事变得简单：

1. 使机器学习专家能够轻松处理地理空间数据；
2. 使遥感专家能够探索机器学习解决方案。

社区：

打包：

测试：

安装

推荐使用 pip 安装 TorchGeo：

pip install torchgeo

有关 conda 和 spack 的安装说明，请参阅 文档。

文档

您可以在 ReadTheDocs 上找到 TorchGeo 的文档。其中包括 API 文档、贡献指南以及多个 教程。更多详细信息，请查看我们的 论文、博客文章 和 YouTube 频道。

示例用法

以下部分提供了一些使用 TorchGeo 的基本示例。

首先，我们导入将在后续章节中使用的各类类和函数：

from lightning.pytorch import Trainer
from torch.utils.data import DataLoader

from torchgeo.datamodules import InriaAerialImageLabelingDataModule
from torchgeo.datasets import CDL, Landsat7, Landsat8, VHR10, stack_samples
from torchgeo.samplers import RandomGeoSampler
from torchgeo.trainers import SemanticSegmentationTask

地理空间数据集与采样器

许多遥感应用都涉及处理地理空间数据集——即带有地理元数据的数据集。由于数据种类繁多，这类数据集往往难以直接使用。地理空间影像通常是多光谱的，不同卫星的光谱波段数量和空间分辨率各不相同。此外，每个文件可能采用不同的坐标参考系统（CRS），因此需要将数据重新投影到一致的坐标系中。

在本示例中，我们展示了如何使用TorchGeo轻松处理地理空间数据，并从Landsat和作物数据层（CDL）数据的组合中采样小块图像。首先，假设用户已下载了Landsat 7和8的影像。由于Landsat 8的光谱波段比Landsat 7多，我们仅使用两颗卫星共有的波段。通过取这两个数据集的并集，我们可以创建一个包含Landsat 7和8所有影像的单一数据集。

landsat7 = Landsat7(paths="...", bands=["B1", ..., "B7"])
landsat8 = Landsat8(paths="...", bands=["B2", ..., "B8"])
landsat = landsat7 | landsat8

接下来，我们对该数据集与CDL数据集取交集。这里使用交集而非并集，是为了确保只从同时拥有Landsat和CDL数据的区域进行采样。需要注意的是，CDL数据可以自动下载并校验其校验和。此外，这些数据集中的文件可能采用不同的坐标参考系统或分辨率，但TorchGeo会自动确保使用匹配的坐标系和分辨率。

cdl = CDL(paths="...", download=True, checksum=True)
dataset = landsat & cdl

现在，这个数据集可以与PyTorch的数据加载器一起使用。与基准数据集不同，地理空间数据集通常包含非常大的图像。例如，CDL数据集由一张覆盖整个美国本土的单幅图像组成。为了能够基于地理坐标从这些数据集中采样，TorchGeo定义了一系列采样器。在本例中，我们使用随机采样器，每次返回256×256像素的图像，并且每轮迭代采样10,000个样本。同时，我们还使用自定义的批处理函数，将每个样本字典合并成一个小批量样本。

sampler = RandomGeoSampler(dataset, size=256, length=10000)
dataloader = DataLoader(dataset, batch_size=128, sampler=sampler, collate_fn=stack_samples)

现在，这个数据加载器就可以用于常规的训练或评估流程中。

for batch in dataloader:
    image = batch["image"]
    mask = batch["mask"]

    # 训练模型，或使用预训练模型进行预测

许多应用都需要根据地理空间元数据智能地组合数据集。例如，用户可能希望：

• 将多个影像源的数据集合并，并将其视为等效（如Landsat 7和8）；
• 将来自不同地理区域的数据集合并（如切萨皮克湾的纽约州和宾夕法尼亚州部分）。

这些组合要求所有查询至少存在于其中一个数据集中，可以通过UnionDataset来实现。类似地，用户也可能希望：

• 将影像和目标标签结合起来，并同时从中采样（如Landsat和CDL）；
• 将多个影像源的数据集合并，用于多模态学习或数据融合（如Landsat和Sentinel）。

这些组合则要求所有查询同时存在于两个数据集中，可以通过IntersectionDataset来实现。当您使用交集（&）和并集（|）运算符时，TorchGeo会自动为您完成这些数据集的组合。

基准数据集

TorchGeo包含一系列基准数据集，这些数据集同时包含输入图像和目标标签。它们涵盖了图像分类、回归、语义分割、目标检测、实例分割、变化检测等多种任务。

如果您之前使用过torchvision，那么这些数据集应该会显得非常熟悉。在本例中，我们将创建西北工业大学（NWPU）的超高分辨率十类目标检测数据集（VHR-10）。该数据集可以像torchvision中的数据集一样，自动下载、校验并解压。

from torch.utils.data import DataLoader

from torchgeo.datamodules.utils import collate_fn_detection
from torchgeo.datasets import VHR10

# 初始化数据集
dataset = VHR10(root="...", download=True, checksum=True)

# 使用自定义批处理函数初始化数据加载器
dataloader = DataLoader(
    dataset,
    batch_size=128,
    shuffle=True,
    num_workers=4,
    collate_fn=collate_fn_detection,
)

# 训练循环
for batch in dataloader:
    image = batch["image"]  # 图像列表
    bbox_xyxy = batch["bbox_xyxy"]  # 边界框列表
    label = batch["label"]  # 标签列表
    mask = batch["mask"]  # 掩码列表

    # 训练模型，或使用预训练模型进行预测

TorchGeo中的所有数据集都与PyTorch的数据加载器兼容，因此很容易集成到现有的训练工作流中。TorchGeo中的基准数据集与torchvision中类似的数据集唯一的区别在于，每个数据集返回的是一个字典，其中包含对应于每个PyTorch张量的键。

预训练权重

预训练权重已被证明对计算机视觉领域的迁移学习任务极为有益。从业者通常会使用在包含 RGB 图像的 ImageNet 数据集上预训练的模型。然而，遥感数据往往不仅限于 RGB，还包含多个多光谱通道，且这些通道的数量和波段范围可能因传感器而异。TorchGeo 是首个支持在不同多光谱传感器上预训练模型的库，并采用了 torchvision 的 多权重 API。当前可用的权重概览可在文档中找到。要创建一个基于 Sentinel-2 影像预训练权重的 timm Resnet-18 模型，可以按如下方式操作：

import timm
from torchgeo.models import ResNet18_Weights

weights = ResNet18_Weights.SENTINEL2_ALL_MOCO
model = timm.create_model("resnet18", in_chans=weights.meta["in_chans"], num_classes=10)
model.load_state_dict(weights.get_state_dict(progress=True), strict=False)

这些权重也可以通过 weights 参数直接用于 TorchGeo Lightning 模块中，具体示例见下文。有关笔记本示例，请参阅此教程。

使用 Lightning 实现可复现性

为了便于将文献中发表的结果与使用 TorchGeo 数据集进行实验的结果直接比较，并进一步减少运行实验所需的样板代码量，我们创建了具有明确训练/验证/测试划分的 Lightning 数据模块 以及适用于分类、回归和语义分割等多种任务的 训练器。这些数据模块展示了如何集成 kornia 库中的数据增强技术、包含预处理变换（并使用预先计算的通道统计信息），同时允许用户轻松地调整与数据本身相关的超参数（而非建模过程中的超参数）。例如，在 Inria 航拍图像标注 数据集上训练一个语义分割模型，只需几行导入语句和四行代码即可完成。

datamodule = InriaAerialImageLabelingDataModule(root="...", batch_size=64, num_workers=6)
task = SemanticSegmentationTask(
    model="unet",
    backbone="resnet50",
    weights=True,
    in_channels=3,
    task="binary",
    loss="bce",
)
trainer = Trainer(default_root_dir="...")
trainer.fit(model=task, datamodule=datamodule)

TorchGeo 还支持使用 LightningCLI 进行命令行训练。可以通过以下两种方式调用：

# 如果已安装 torchgeo
torchgeo
# 如果已安装 torchgeo，或已将其克隆到当前目录
python3 -m torchgeo

它支持通过命令行配置或 YAML/JSON 配置文件进行设置。有效选项可通过帮助信息查看：

# 查看有效阶段
torchgeo --help
# 查看有效训练器选项
torchgeo fit --help
# 查看有效模型选项
torchgeo fit --model.help ClassificationTask
# 查看有效数据选项
torchgeo fit --data.help EuroSAT100DataModule

使用以下配置文件：

trainer:
  max_epochs: 20
model:
  class_path: ClassificationTask
  init_args:
    model: 'resnet18'
    in_channels: 13
    num_classes: 10
data:
  class_path: EuroSAT100DataModule
  init_args:
    batch_size: 8
  dict_kwargs:
    download: true

我们可以看到该脚本的实际运行效果：

# 训练并验证模型
torchgeo fit --config config.yaml
# 仅验证
torchgeo validate --config config.yaml --ckpt_path=...
# 计算并报告测试准确率
torchgeo test --config config.yaml --ckpt_path=...

如果需要扩展功能，也可以将其导入 Python 脚本中使用：

from torchgeo.main import main

main(["fit", "--config", "config.yaml"])

更多详细信息请参阅 Lightning 文档。

引用

如果您在工作中使用了本软件，请引用我们的论文：

@article{stewart2024torchgeo,
    author = {Stewart, Adam J. and Robinson, Caleb and Corley, Isaac A. and Ortiz, Anthony and Lavista Ferres, Juan M. and Banerjee, Arindam},
    title = {{TorchGeo}: Deep Learning With Geospatial Data},
    journal = {ACM Transactions on Spatial Algorithms and Systems},
    volume = {11},
    number = {4},
    pages = {1--28},
    month = aug,
    year = {2025}
}

贡献

本项目欢迎各类贡献和建议。如果您希望提交拉取请求，请参阅我们的贡献指南以获取更多信息。

本项目已采纳 Contributor Covenant 行为准则。如需更多信息，请参阅 Contributor Covenant 行为准则常见问题解答，或通过 Slack 联系 @adamjstewart，提出任何其他问题或意见。

TorchGeo 快速上手指南

TorchGeo 是一个基于 PyTorch 的地理空间数据领域库，类似于 torchvision。它提供了专门针对地理空间数据的 datasets（数据集）、samplers（采样器）、transforms（变换）和预训练模型，旨在降低机器学习专家处理地理空间数据的门槛，同时帮助遥感专家探索机器学习解决方案。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux, macOS 或 Windows
• Python: 3.8 或更高版本
• 核心依赖:
  • PyTorch (建议安装与您的 CUDA 版本匹配的最新版)
  • Lightning (可选，用于复现实验和简化训练流程)
• 地理空间依赖: TorchGeo 底层依赖 rasterio, gdal, shapely 等库。
  • 推荐方案: 为了避免编译 GDAL 等 C++ 库的复杂性，强烈建议使用 Conda 环境进行管理。

安装步骤

方式一：使用 pip 安装（推荐）

如果您已经配置好了包含 GDAL 等地理空间库的环境，可以直接使用 pip 安装：

pip install torchgeo

国内加速提示: 如果下载速度较慢，可以使用清华或阿里镜像源：

pip install torchgeo -i https://pypi.tuna.tsinghua.edu.cn/simple

方式二：使用 Conda 安装（最稳健）

此方式会自动解决复杂的地理空间依赖（如 GDAL、PROJ 等），适合大多数用户：

conda install -c conda-forge torchgeo

国内加速提示: 配置 conda-forge 国内镜像后速度更快：

conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge/
conda config --set show_channel_urls yes
conda install -c conda-forge torchgeo

基本使用

TorchGeo 的核心优势在于能够轻松组合不同的地理空间数据集，并自动处理坐标参考系统（CRS）和分辨率的对齐。

以下示例展示了如何组合 Landsat 卫星影像数据和 CDL (Cropland Data Layer) 土地覆盖数据，并使用采样器生成训练批次。

1. 导入必要模块

from torch.utils.data import DataLoader

from torchgeo.datasets import CDL, Landsat7, Landsat8, stack_samples
from torchgeo.samplers import RandomGeoSampler

2. 构建组合数据集

TorchGeo 允许使用集合运算符来组合数据集：

• | (Union): 取并集，适用于合并不同来源但类型相同的数据（如 Landsat 7 和 8）。
• & (Intersection): 取交集，适用于将影像与标签对齐（确保只在两者都覆盖的区域采样）。

# 初始化 Landsat 7 和 8 数据集 (假设路径已准备好)
# 仅选择两颗卫星共有的波段
landsat7 = Landsat7(paths="path/to/landsat7", bands=["B1", "B2", "B3", "B4", "B5", "B6", "B7"])
landsat8 = Landsat8(paths="path/to/landsat8", bands=["B2", "B3", "B4", "B5", "B6", "B7", "B8"])

# 合并 Landsat 7 和 8 (并集)
landsat = landsat7 | landsat8

# 初始化 CDL 数据集 (支持自动下载和校验)
cdl = CDL(paths="path/to/cdl", download=True, checksum=True)

# 获取影像和标签的重叠区域 (交集)
# TorchGeo 会自动处理重投影和分辨率匹配
dataset = landsat & cdl

3. 定义采样器与数据加载器

地理空间图像通常非常大（例如覆盖整个美国），不能直接一次性加载。需要使用 GeoSampler 按坐标块进行采样。

# 定义随机采样器：每次采集 256x256 像素的图块，每个 epoch 采样 10000 次
sampler = RandomGeoSampler(dataset, size=256, length=10000)

# 创建 DataLoader
# collate_fn=stack_samples 用于将样本字典堆叠成 batch
dataloader = DataLoader(
    dataset, 
    batch_size=32, 
    sampler=sampler, 
    collate_fn=stack_samples,
    num_workers=4
)

4. 开始训练循环

现在可以像使用普通 PyTorch 数据一样进行迭代：

for batch in dataloader:
    # batch 是一个字典，包含 'image', 'mask' 等键
    image = batch["image"]  # 形状: [Batch, Channels, Height, Width]
    mask = batch["mask"]    # 形状: [Batch, Classes, Height, Width]
    
    # 在此处插入您的模型训练代码
    # outputs = model(image)
    # loss = criterion(outputs, mask)
    pass

进阶：使用预训练模型

TorchGeo 支持针对多光谱传感器（如 Sentinel-2）预训练的权重，这对于迁移学习非常有用：

import timm
from torchgeo.models import ResNet18_Weights

# 加载在 Sentinel-2 数据上预训练的权重
weights = ResNet18_Weights.SENTINEL2_ALL_MOCO

# 创建模型，注意输入通道数需匹配预训练权重
model = timm.create_model("resnet18", in_chans=weights.meta["in_chans"], num_classes=10)

# 载入权重
model.load_state_dict(weights.get_state_dict(progress=True), strict=False)