torchgeo

TorchGeo 项目介绍

简介

TorchGeo 是一个专为地理空间数据设计的 PyTorch 域库，与 PyTorch 的 torchvision 类似。它提供了数据集、采样器、变换工具以及预训练模型，旨在帮助机器学习专家和遥感专家更简单地处理和探索地理空间数据的机器学习解决方案。

项目目标

TorchGeo 的目标很明确：

为机器学习专家简化地理空间数据的使用。
帮助遥感专家更轻松地探索机器学习解决方案。

安装方法

安装 TorchGeo 的推荐方式是使用 pip，可以通过以下命令进行安装：

$ pip install torchgeo

如需了解如何使用 conda 或 spack 进行安装，请参阅文档。

项目功能

地理空间数据集和采样器

许多遥感应用涉及与地理空间数据集合作，这些数据集带有地理元数据，处理这些数据集可能会比较复杂，因为地理空间图像通常是多光谱的，每颗卫星都有不同的光谱波段和空间分辨率。此外，文件可能采用不同的坐标参考系统，这需要对数据进行重投影。

TorchGeo 使得处理 Landsat 和农作物数据图层（CDL）等数据变得非常简单，可以从中采样小图像块。甚至可以将 Landsat 7 和 8 数据集中的图像综合成一个数据集。

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

基准数据集

TorchGeo 包含多个基准数据集，这些数据集包括输入图像和目标标签。涉及图像分类、回归、语义分割、目标检测等多种任务。

from torchgeo.datasets import VHR10

dataset = VHR10(root="...", download=True, checksum=True)

预训练权重

在计算机视觉领域，预训练权重对于迁移学习任务具有重要价值。TorchGeo 支持不同多光谱传感器上预训练的模型，并采用 torchvision 的多权重 API，比如可以加载在 Sentinel-2 图像上预训练的 ResNet18 模型。

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)

重现性与 Lightning

为了便于与文献中发表的结果进行比较，并减少运行实验所需的样板代码，TorchGeo 提供了 Lightning 数据模块和训练器，这些工具让用户能够轻松实验数据相关的超参数，简化语义分割模型在如 Inria Aerial Image Labeling 数据集上的训练流程。

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