案例研究

IBM 研究院如何利用 PyTorch 和 TerraTorch 让每个人都能使用地理空间计算机视觉

作者 2025年5月1日2025年8月19日暂无评论

基于地球观测的分析对于了解我们的星球正变得至关重要——从监测森林砍伐到跟踪城市发展,再到分析气候变化的影响。然而,对于许多从业者来说,将人工智能模型应用于卫星图像和地球观测数据所需的编码及深度学习技能,一直是巨大的障碍。

随着 IBM Research 发布 TerraTorch 1.0(一个用于微调地理空间计算机视觉基础模型的 PyTorch 领域库),我们不仅让地理空间人工智能变得更易于使用,也使其对更广泛的 PyTorch 社区而言更具实用性。我们的目标是:简化流程,让任何数据科学家、研究人员或爱好者都能轻松构建强大的地理空间模型,且对 GPU 和数据处理的需求较低。

 

基础模型的力量在于,即使去除了 75-95% 的输入数据,模型在重建输入数据方面依然表现出色——因此能在深层的潜在空间中学习我们星球底层的物理规律。

商业挑战

我们的目标是消除阻碍人们大规模处理卫星图像、天气和气候数据的技术壁垒。我们与 NASA 合作开发了 Prithvi 基础模型系列。利用 PyTorch 提供的简洁 API 集成 AI 研究的最新创新成果,极大地促进了这项工作。

我们希望创建一个框架,让任何人只需几步操作即可从原始数据过渡到可供推理的模型。

 

一个在 PyTorch 上创建并微调的天气和气候基础模型如何用于天气预报

IBM Research 如何使用 PyTorch

我们在 PyTorch 之上构建了 TerraTorch,利用其动态生态系统集成了:

  • 用于简洁、可扩展训练循环的 PyTorch Lightning
  • 用于地理空间数据处理和转换(PyTorch 变换)的 TorchGeo
  • 对于诸如由 IBM 与 ESA 共同开发的领先生成式多模态基础模型“Terramind”,以及由 IBM 与 NASA 共同开发的“Prithvi”系列等基础模型,TerraTorch 已被用于微调所有下游地理空间卫星图像、天气和气候数据模型。它包含了 IBM 作为 Granite 一部分发布的一系列微调模型。此外,其他引人注目的基础模型和生态系统组件,如 Clay、SatMAE、Satlas、DeCur 和 DOFA,也被纳入了 TerraTorch 中。
  • 强大且最先进的视觉 Transformer,用于试验现代神经网络架构
  • TerraTorch-Iterate 构建于 PyTorch、Optuna、MLFlow 和 Ray Tune 之上,用于超参数优化 (HPO)、神经架构搜索 (NAS) 和基础模型基准测试 (GeoBench),其中 TerraTorch 成为了参考实现。

微调和推理过程完全描述在一个 YAML 配置文件中。在该文件中,定义了模型的架构构建模块(骨干网络、颈部、解码器、头部)。模型工厂(Model Factory)使用内置和自定义注册表来组装模型。此外,优化器(Optimizer)和数据模块(Data Modules)也是按照配置文件的定义进行创建。最后,所有内容都被传递给 Lightning Trainer 执行任务。

得益于 PyTorch 的灵活性,我们能够快速进行原型设计、迭代模型架构并为一系列地理空间应用部署流水线——从洪水和生物量检测到提高气候数据分辨率,其中我们的一些工作已成为 IBM Granite 地理空间模型系列的一部分。

 

IBM Research 与 NASA 共同开发的 Prithvi-EO-2.0-600M 基础模型架构

用 PyTorch 解决 AI 挑战

PyTorch 帮助我们解决了三大挑战:

  • 易于实验:动态计算图、自动微分、CUDA 的完全抽象以及丰富的可视化工具,使得测试不同的模型和训练策略变得简单。
  • 可扩展性:通过 DDP、FSDP、PyTorch Lightning 和 TorchGeo,我们可以在大规模数据集上训练模型,而无需担心基础设施问题。
  • 社区支持:PyTorch 作为 AI 研究的事实标准,拥有活跃的社区和优秀的文档,使我们能够轻松克服障碍,并紧跟 AI 研究的最新进展。

来自 IBM Research 的寄语

“PyTorch 让我有能力将复杂的线性代数和优化问题转化为社区可访问、可共享的解决方案。能够为任何通过 AI 了解我们星球的好奇者构建和微调模型,这让我感到充满力量。”

— Romeo Kienzler,IBM Research 苏黎世(吕什利孔)AI 研究工程师

使用 PyTorch 的好处

使用 PyTorch 使我们能够:

  • 构建一个可复现的开源框架,用于微调地理空间基础模型
  • 通过易于遵循的 Notebook、TerraTorch 配置文件、教程以及 HuggingFace 上的模型检查点与社区分享我们的工作
  • 快速迭代基础模型架构并部署微调模型以进行推理,从而实现从研究到现实世界客户产品的转化

了解更多

有关该项目的更多信息并探索代码,请访问: