基于地球观测的分析对于理解我们的星球正变得至关重要——从监测森林砍伐到跟踪城市发展,再到分析气候变化的影响。然而,将人工智能模型应用于卫星影像和地球观测数据的编程和深度学习技能,传统上一直是许多从业者的主要障碍。
随着 IBM 研究院发布 TerraTorch 1.0——一个用于微调地理空间计算机视觉基础模型的 PyTorch 领域库,我们不仅让地理空间人工智能变得更容易获取,也使其对更广泛的 PyTorch 社区更具实用性。我们的目标是:简化流程,以便任何数据科学家、研究人员或爱好者都能以较低的 GPU 和数据处理要求,轻松构建强大的地理空间模型。
基础模型的强大之处在于,即使移除了 75-95% 的输入数据,模型在重构输入数据方面仍然表现出色——从而在一个深层的潜在空间中学习我们星球的潜在物理规律。
商业挑战
我们的目标是消除那些阻碍人们大规模处理卫星影像、天气和气候数据的技术壁垒。我们与 NASA 合作,共同开发了 Prithvi 系列基础模型。利用 PyTorch 提供的简洁 API,整合人工智能研究的最新创新,使这项工作变得更加容易。
我们希望创建一个任何人都可以使用的框架,只需几个步骤就能从原始数据生成可用于推理的模型。
一个在 PyTorch 上创建和微调的天气与气候基础模型如何用于天气预报
IBM 研究院如何使用 PyTorch
我们基于 PyTorch 构建了 TerraTorch,利用其充满活力的生态系统来集成:
- PyTorch Lightning,用于实现简洁、可扩展的训练循环。
- TorchGeo,用于处理地理空间数据和转换(PyTorch 变换)。
- 对于像 IBM 和 ESA 共同开发的领先生成式多模态基础模型“Terramind”,以及 IBM 和 NASA 共同开发的“Prithvi”系列这样的基础模型,TerraTorch 已被用于微调所有下游的地理空间模型,这些模型适用于卫星影像、天气和气候数据。这包括 IBM 作为Granite 的一部分发布的一系列微调模型。此外,TerraTorch 还集成了其他有趣的基础模型和生态系统组件,如 Clay、SatMAE、Satlas、DeCur 和 DOFA。
- 强大且先进的视觉 Transformer,用于实验现代神经网络架构。
- TerraTorch-Iterate 构建在 PyTorch、Optuna、MLFlow 和 Ray Tune 之上,用于超参数优化(HPO)、神经架构搜索(NAS)和基础模型基准测试(GeoBench),其中 TerraTorch 成为了参考实现。
微调和推理过程完全通过一个 YAML 配置文件来描述。在该文件中,定义了模型的架构构建块(主干、颈部、解码器、头部)。模型工厂使用内置和自定义的注册表来组装模型。此外,优化器和数据模块也根据配置文件的定义被创建。最后,所有这些都传递给 Lightning Trainer,由它来执行任务。
借助 PyTorch 的灵活性,我们能够快速进行原型设计,迭代模型架构,并为一系列地理空间应用部署管道——从洪水和生物量检测到提高气候数据分辨率,我们的一些工作已成为IBM Granite 地理空间模型系列的一部分。
IBM 研究院与 NASA 共同开发的 Prithvi-EO-2.0-600M 基础模型架构
用 PyTorch 解决人工智能挑战
PyTorch 帮助我们应对了三大挑战:
- 易于实验:动态计算图、自动微分、对 CUDA 的完全抽象以及丰富的可视化工具,使测试不同模型和训练策略变得简单。
- 可扩展性:借助 DDP、FSDP、PyTorch Lightning 和 TorchGeo,我们可以在大规模数据集上训练模型,而无需担心基础设施问题。
- 社区支持:PyTorch 作为人工智能研究领域的事实标准,其活跃的社区和优秀的文档使我们能够轻松克服障碍,并与人工智能研究的最新进展保持同步。
来自 IBM 研究院的寄语
“PyTorch 让我有能力将复杂的线性代数和优化问题转化为社区可以访问和共享的解决方案。我们正在为任何对通过人工智能理解我们星球感到好奇的人构建和微调模型,这让我感到很有力量。”
— Romeo Kienzler, IBM 苏黎世研究院(Rueschlikon)人工智能研究工程师
使用 PyTorch 的好处
使用 PyTorch 使我们能够:
- 构建一个可复现的、开源的框架,用于微调地理空间基础模型。
- 通过易于理解的笔记本、TerraTorch 配置文件、教程和 HuggingFace 上的模型检查点,与社区分享我们的工作。
- 快速迭代基础模型架构,并部署微调模型用于推理,从研究到实际的客户产品。
了解更多
有关此项目的更多信息和代码,请访问: