NVIDIA/physicsnemo

Open-source deep-learning framework for building, training, and fine-tuning deep learning models using state-of-the-art Physics-ML methods

解決的問題

PhysicsNeMo 是一個旨在加速科學與工程領域 AI 模型（AI4Science）開發的深度學習框架。它解決了將物理知識與數據結合以創建複雜物理系統高保真、即時預測的挑戰，縮短了計算流體力學 (CFD)、結構力學和電磁學等領域模擬開發所需的時間。

工作原理

該框架構建在 PyTorch 之上，提供了一套模組化的組件堆疊，用於構建、訓練和部署物理資訊機器學習模型。它包括：

  • Model Zoo：包含神經算子 (FNO, DeepONet)、圖神經網絡 (GNNs)、擴散模型和物理資訊神經網絡 (PINNs) 在內的優化架構集合。
  • Data Pipelines：專門針對網格和點雲等科學數據結構進行調優的可擴展數據管道。
  • Distributed Computing：基於 torch.distributed 的子模組，用於將訓練從單 GPU 擴展到多節點集群。
  • Symbolic PDE Utilities：透過 SymPy 定義方程式並使用自動空間導數計算物理資訊損失的工具。

適用對象

面向需要在 NVIDIA GPU 上以最佳性能訓練大規模物理-ML 模型的 SciML 研究人員、氣候科學家和工程領域專家。

亮點

  • 無縫 PyTorch 集成：作為 PyTorch 的補充工具運行，具有熟悉的 Python 介面。
  • GPU 優化：針對 NVIDIA GPU 進行了高度優化，以實現最大的訓練速度和可擴展性。
  • 可擴展設計：支持 ONNX 以實現靈活部署，並提供用於定義自定義物理模型和約束的 API。
  • 特定領域套件：包括用於 CFD 和天氣/氣候研究 (Earth-2 Studio) 的專業子模組。