Hugging Face
PyPI版本
测试状态
文档支持

TerraTorch是什么？

TerraTorch 是一个基于 PyTorch Lightning 和地理空间数据处理库 TorchGeo 的开源工具库。它的核心目标是为地理空间基础模型（Geospatial Foundation Models）提供高度灵活的微调框架，支持从命令行接口（CLI）到Jupyter笔记本的多层次交互方式。

为什么选择TerraTorch？

  • 灵活的任务支持：覆盖图像分割（Image Segmentation）、分类（Classification）和像素级回归（Pixel-wise Regression）等任务。
  • 预训练模型集成：无缝对接Prithvi、SatMAE、ScaleMAE等地理空间领域主流模型。
  • 模块化设计：通过模型工厂（Model Factories）快速组合不同主干网络（Backbones）与解码器（Decoders）。
  • 数据集即插即用：内置GEO-Bench和TorchGeo数据集，无需额外编写数据加载代码。

核心功能解析

1. 模型工具：从训练到部署

TerraTorch提供以下关键建模工具：

  • 定制化训练器（Trainers）：针对不同任务优化的训练流程，支持分布式训练与混合精度计算。
  • 模型工厂模式：例如，将Prithvi-100M主干网络与Segmentation Models PyTorch的解码器结合，快速构建分割模型。
  • 数据集自动化：通过数据模块（Datamodules）直接读取地理空间数据格式（如GeoTIFF），无需预处理代码。

2. 预训练模型库

TerraTorch整合了六大类地理空间预训练模型：

  1. Prithvi系列：IBM与NASA联合开发的100M参数多模态模型。
  2. SatMAE与ScaleMAE：基于自监督学习的遥感图像编码器。
  3. Satlas与DOFA：来自TorchGeo的高性能预训练权重。
  4. SSL4EO系列：适用于多光谱与合成孔径雷达（SAR）数据。
  5. Clay模型：专注于地质与气候预测任务。
  6. Timm与MMSegmentation生态：支持超过200种通用视觉模型。

3. 数据集支持

  • GEO-Bench：涵盖土地利用分类、森林监测等场景的标准化数据集。
  • TorchGeo数据集：包括EuroSAT、NAIP、LandCover.ai等公开数据源。
  • 自定义数据兼容性：用户仅需指定数据目录路径即可完成加载。

安装与配置指南

环境要求

  • Python版本：建议3.8及以上
  • 依赖管理：需预先安装GDAL库。推荐通过Conda简化安装：

    conda install -c conda-forge gdal

快速安装

  • 稳定版：

    pip install terratorch==<版本号>
  • 开发版（获取最新功能）：

    pip install git+https://github.com/IBM/terratorch.git

开发者模式

若需贡献代码或扩展功能，按以下步骤配置：

git clone https://github.com/IBM/terratorch.git
cd terratorch
pip install -r requirements_test.txt
conda install -c conda-forge gdal
pip install -e .

气象模型支持（实验性）

对气象任务（如天气预报）感兴趣的用户，可通过以下命令启用扩展：

pip install -e .[wxc]

注：当前仅支持Python 3.11及以上版本。

实战案例：从零微调地理空间模型

步骤1：选择任务与数据集

假设需进行土地覆盖分类（Land Cover Classification），选择EuroSAT数据集：

from terratorch.datasets import EuroSATDataModule
datamodule = EuroSATDataModule(data_root="/path/to/eurosat")

步骤2：构建模型

使用Prithvi主干网络与线性分类头：

from terratorch.models import create_model
model = create_model(
    backbone="prithvi_100m",
    task="classification",
    num_classes=10
)

步骤3：启动训练

通过PyTorch Lightning Trainer接口配置训练参数：

from pytorch_lightning import Trainer
trainer = Trainer(max_epochs=50, devices=4, accelerator="gpu")
trainer.fit(model, datamodule=datamodule)

文档与社区资源

官方文档

如何参与贡献？

TerraTorch遵循开源协作模式，欢迎通过以下方式参与：

  1. 提交问题：在GitHub仓库报告Bug或提出功能建议。
  2. 代码贡献：遵循开发者安装指南，签署提交（Sign-off Commits）：

    git commit -s -m "修复图像加载异常"
  3. 模型共享：将微调后的模型权重发布至Hugging Face Hub。

总结

TerraTorch通过其高度模块化的设计，降低了地理空间AI模型的应用门槛。无论是研究者还是工程师，均可快速实现从数据加载、模型微调到部署的全流程。其与TorchGeo、Hugging Face生态的深度集成，进一步拓展了在遥感、气象、环境监测等领域的应用场景。

如需了解更多技术细节，请访问官方GitHub仓库或加入Hugging Face社区讨论。