NVIDIA PhysicsNeMo

📝 NVIDIA PhysicsNeMo 正在升级至 v2.0——所有功能均得以保留，同时安装和与外部包的集成更加简便。更多详情请参阅迁移指南！

NVIDIA PhysicsNeMo | 文档 | 安装指南 | 快速入门 | 贡献指南 | 开发博客

什么是 PhysicsNeMo？

NVIDIA PhysicsNeMo 是一个开源深度学习框架，用于构建、训练、微调和推理物理人工智能模型，采用最先进的 SciML 方法，服务于 AI4Science 和工程领域。

PhysicsNeMo 提供 Python 模块，可组合出可扩展且优化的训练与推理流水线，以探索、开发、验证和部署结合物理知识与数据的人工智能模型，从而实现实时预测。

无论您是在探索神经算子、图神经网络或 Transformer 模型，还是对物理信息神经网络或介于两者之间的混合方法感兴趣，PhysicsNeMo 都能为您提供一套优化的工具栈，帮助您大规模训练模型。

• 关于 PhysicsNeMo 的更多信息
  • 可扩展的 GPU 优化训练库
  • 一系列物理信息机器学习模型
  • 无缝的 PyTorch 集成
  • 易于自定义和扩展
  • AI4Science 库
    • 领域专用包
• 谁在使用和贡献 PhysicsNeMo
• 为什么使用 PhysicsNeMo
• 快速入门
• 资源
• 安装
• 贡献
• 交流
• 许可证

关于 PhysicsNeMo 的更多信息

从细粒度来看，PhysicsNeMo 被设计为模块化功能，因此提供了内置的可组合模块，并将其封装成几个关键组件：

组件	描述
physicsnemo.models（更多细节）	一系列优化、可定制且易于使用的模型架构家族，如神经算子、图神经网络、扩散模型、Transformer 模型等
physicsnemo.datapipes	优化且可扩展的内置数据管道，专为处理点云、网格等工程和科学数据结构而设计
physicsnemo.distributed	基于 torch.distributed 构建的分布式计算子模块，只需几步即可实现并行训练
physicsnemo.curator	一个子模块，用于简化和加速工程数据集的数据整理流程
physicsnemo.sym.geometry	一个子模块，用于处理基于构造实体几何建模和 STL 格式 CAD 文件的几何数据，以支持深度学习训练
physicsnemo.sym.eq	一个子模块，允许在深度学习训练中使用偏微分方程，提供多种常见方程的实现方式，并便于用户自定义

完整列表请参阅 PhysicsNeMo API 文档中的 PhysicsNeMo 部分。

AI4Science 库

通常，PhysicsNeMo 可用作以下两种场景之一：

• 作为 PyTorch 的补充工具，在探索 SciML 和 AI4Science 应用时使用。
• 作为深度学习研究平台，在 NVIDIA GPU 上提供规模性和最佳性能。

领域专用包

以下是专为特定领域的专家设计的软件包，以满足其独特的研究需求：

• PhysicsNeMo CFD：PhysicsNeMo 的推理子模块，旨在帮助 CFD 领域专家利用预训练的人工智能模型进行探索、实验和验证。
• PhysicsNeMo Curator：PhysicsNeMo 的推理子模块，用于简化和加速工程数据集的数据整理过程。
• Earth-2 Studio：PhysicsNeMo 的推理子模块，使气候研究人员能够利用预训练的人工智能模型探索和试验天气及气候相关问题。

可扩展的 GPU 优化训练库

PhysicsNeMo 提供了一个高度优化且可扩展的训练库，旨在充分发挥 NVIDIA GPU 的强大性能。 分布式计算 工具允许用户仅用几行代码即可高效地从单个 GPU 扩展到多节点 GPU 集群，从而确保大规模物理信息驱动的机器学习（ML）模型能够快速有效地进行训练。 该框架还包含对高级 优化工具、量身定制的 数据管道 以及 验证工具 的支持，以提升端到端训练速度。

一系列物理信息驱动的 ML 模型

PhysicsNeMo 提供了一个专为物理信息驱动的机器学习应用设计的先进模型库。用户可以使用底层 PyTorch 层，并将其与精选的 PhysicsNeMo 层相结合，构建任意模型架构。

模型动物园 包含了多种模型家族的优化实现，例如神经算子：

• 傅里叶神经算子（FNOs）
• DeepONet
• DoMINO
• 图神经网络（GNNs）
• MeshGraphNet
• 用于拉格朗日方法的 MeshGraphNet
• XAeroNet
• 扩散模型
• 校正扩散模型
• DDPM
• PhysicsNeMo GraphCast
• Transsolver
• RNNs
• SwinVRNN
• 物理信息神经网络（PINNs）

以及其他众多模型。

这些模型针对不同的物理领域进行了优化，例如计算流体动力学、结构力学和电磁学。用户可以下载、自定义并在此基础上构建模型，以满足其特定需求，从而显著缩短开发高保真度仿真所需的时间。

与 PyTorch 的无缝集成

PhysicsNeMo 构建在 PyTorch 之上，为已经熟悉 PyTorch 的用户提供了一种熟悉且友好的使用体验。 这包括简单的 Python 接口和模块化设计，使得用户能够轻松地将 PhysicsNeMo 与现有的 PyTorch 工作流程结合使用。 用户可以充分利用 PyTorch 庞大的生态系统及其丰富的库和工具，同时受益于 PhysicsNeMo 在物理信息驱动机器学习方面的专业能力。这种无缝集成确保用户无需经历陡峭的学习曲线即可快速上手使用 PhysicsNeMo。

更多信息请参阅 将 PyTorch 模型转换为 PhysicsNeMo 模型。

易于自定义和扩展

PhysicsNeMo 设计为高度可扩展，允许用户以最小的努力添加新功能。该框架提供了 Python 式的 API，用于定义新的物理模型、几何形状和约束条件，从而轻松扩展其功能以适应新的应用场景。 PhysicsNeMo 的适应性进一步通过以下关键特性得到增强：支持 ONNX 格式，便于灵活部署模型；强大的 日志记录工具，可简化错误处理；以及高效的 检查点机制，便于模型的加载和保存。

这种可扩展性确保了 PhysicsNeMo 能够适应研究人员和工程师不断变化的需求，从而促进物理信息驱动机器学习领域的创新解决方案的开发。

有关功能和特性的详细信息，请参阅 PhysicsNeMo 文档。

参考示例 涵盖了广泛的物理约束和数据驱动的工作流程，以适应科学和工程领域的多样化应用场景。

[!提示] 想了解 PhysicsNeMo 如何帮助您吗？请尝试我们的 [实验性] 聊天机器人——PhysicsNeMo 指南，获取答案。

你好，世界

你可以在自己的 PyTorch 代码中像下面这样简单地开始使用 PhysicsNeMo：

>>> import torch
>>> from physicsnemo.models.mlp.fully_connected import FullyConnected
>>> model = FullyConnected(in_features=32, out_features=64)
>>> input = torch.randn(128, 32)
>>> output = model(input)
>>> output.shape
torch.Size([128, 64])

要使用分布式模块，可以按照以下方式操作（以下示例为分布式数据并行训练；如需更深入的教程，请参阅 PhysicsNeMo 分布式）：

import torch
from torch.nn.parallel import DistributedDataParallel
from physicsnemo.distributed import DistributedManager
from physicsnemo.models.mlp.fully_connected import FullyConnected

def main():
    DistributedManager.initialize()
    dist = DistributedManager()

    arch = FullyConnected(in_features=32, out_features=64).to(dist.device)

    if dist.distributed:
        ddps = torch.cuda.Stream()
        with torch.cuda.stream(ddps):
            arch = DistributedDataParallel(
                arch,
                device_ids=[dist.local_rank],
                output_device=dist.device,
                broadcast_buffers=dist.broadcast_buffers,
                find_unused_parameters=dist.find_unused_parameters,
            )
        torch.cuda.current_stream().wait_stream(ddps)

    # 设置优化器
    optimizer = torch.optim.Adam(
        arch.parameters(),
        lr=0.001,
    )

    def training_step(invar, target):
        pred = arch(invar)
        loss = torch.sum(torch.pow(pred - target, 2))
        loss.backward()
        optimizer.step()
        return loss

    # 示例训练循环
    for i in range(20):
        # 为了简单起见，随机生成输入和目标
        input = torch.randn(128, 32, device=dist.device)
        target = torch.randn(128, 64, device=dist.device)

        # 训练步骤
        loss = training_step(input, target)

if __name__ == "__main__":
    main()

要使用 PDE 模块，可以按照以下方式操作：

>>> from physicsnemo.sym.eq.pdes.navier_stokes import NavierStokes
>>> ns = NavierStokes(nu=0.01, rho=1, dim=2)
>>> ns.pprint()
连续性方程: u__x + v__y
动量方程_x: u*u__x + v*u__y + p__x + u__t - 0.01*u__x__x - 0.01*u__y__y
动量方程_y: u*v__x + v*v__y + p__y + v__t - 0.01*v__x__x - 0.01*v__y__y

谁在使用和贡献于 PhysicsNeMo

PhysicsNeMo 是一个开源项目，得到了 SciML 和 AI4Science 领域研究人员的贡献。虽然 PhysicsNeMo 团队致力于优化底层软件栈，但社区也在协作并贡献模型架构、数据集和参考应用，以便我们能够在开发可泛化的模型架构和算法方面不断创新。

近期社区贡献者的一些例子包括：HP 实验室 3D 打印团队、斯坦福心血管研究团队、伊利诺伊大学厄巴纳-香槟分校团队、卡内基梅隆大学团队等。

近期使用 PhysicsNeMo 的研究团队的例子有：橡树岭国家实验室团队、慕尼黑工业大学 CFD 团队等。

请访问此页面以获取利用 PhysicsNeMo 的完整研究工作列表。有关使用 PhysicsNeMo 的企业的列表，请参阅 PhysicsNeMo 官网。

您是否正在使用 PhysicsNeMo，并希望在 NVIDIA 博客 上展示您的工作？请填写此 提案表单，我们将尽快与您联系！

他们为何使用 PhysicsNeMo

以下是 PhysicsNeMo 在 SciML 模型开发中的主要优势：

SciML 基准测试与验证	轻松使用通用 SciML 模板处理异构数据集	开箱即用的性能与可扩展性
PhysicsNeMo 使研究人员能够针对标准基准问题，使用经过验证的架构，并结合详细的领域特定验证标准来评估其 AI 模型。	PhysicsNeMo 使研究人员能够从最先进的 SciML 架构中选择，并为其用例使用内置的数据管道。	PhysicsNeMo 提供开箱即用的高效训练流水线，包括针对异构工程和科学数据集的优化 ETL 流水线，以及跨多 GPU 和多节点 GPU 的开箱即用扩展能力。

敬请期待同行 SciML 研究人员对 PhysicsNeMo 的评价（即将发布）。

开始使用 PhysicsNeMo

以下资源将帮助您学习如何使用 PhysicsNeMo。最好的方法是从一个参考示例开始，然后根据您自己的用例进行调整。

• 将 PhysicsNeMo 与您的 PyTorch 模型一起使用
• 使用 PhysicsNeMo 内置模型
• 入门指南
• 参考示例
• 用户指南文档

学习 AI 物理学

• 探索 Hugging Face 上的 Jupyter 笔记本
• AI4Science PhysicsNeMo 训练营
• DLI 自定进度培训
• 面向科学与工程的深度学习讲座系列
• 视频教程

资源

• 入门网络研讨会
• PhysicsNeMo：目的与使用方法
• AI4Science PhysicsNeMo训练营
• PhysicsNeMo预训练模型
• PhysicsNeMo数据集及补充材料

安装

您可以通过两种受支持的方式安装 PhysicsNeMo：通过 pip（原生 pip 或 uv）或使用 NVIDIA 容器镜像。请选择适合您的环境和工作流的方法。

以下说明涵盖了 PhysicsNeMo 的基础模块。可选依赖项列在 pyproject.toml 中。训练示例 不包含在 pip 轮子或容器中；请克隆仓库，并将这些示例作为起点。许多示例都包含 requirements.txt 文件，用于安装额外的依赖项。

CUDA 后端选择

重要提示： 若要获得 GPU 加速的 RAPIDS 包（cuML、pylibraft、cupy）以及与 CUDA 版本匹配的 PyTorch 构建，您必须在安装时包含 cu13 或 cu12。诸如 nn-extras 和 utils-extras 等功能扩展提供了额外的非 CUDA 包（scipy、natten、wandb 等），但它们本身并不包含 RAPIDS 依赖项。

PhysicsNeMo 同时支持 CUDA 12 和 CUDA 13 后端。后端是通过一个与功能扩展无关的额外选项来选择的——您可以自由组合：

额外选项	提供的内容
cu13	PyTorch（CUDA 13.0）、cuML-cu13、pylibraft-cu13、cupy-cuda13x
cu12	PyTorch（CUDA 12.8）、cuML-cu12、pylibraft-cu12、cupy-cuda12x
(无)	来自 PyPI 的 PyTorch（默认构建），不含 RAPIDS 包

PyPI

从 PyPI 安装最新版本：

pip install nvidia-physicsnemo
python -c "import physicsnemo; print('PhysicsNeMo 版本:', physicsnemo.__version__)"

要使用特定的 CUDA 后端和可选的功能扩展进行安装：

# CUDA 13 后端，包含 nn-extras
pip install "nvidia-physicsnemo[cu13,nn-extras]"

# CUDA 12 后端，包含 nn-extras
pip install "nvidia-physicsnemo[cu12,nn-extras]"

其他功能扩展（utils-extras、mesh-extras、model-extras、datapipes-extras、gnns）也可以以相同的方式组合使用。

您还可以通过运行 Hello World 示例来验证安装是否成功。

uv

对于开发或直接从源代码运行示例，您可以使用 uv 克隆仓库并同步依赖项：

git clone https://github.com/NVIDIA/physicsnemo.git
cd physicsnemo
uv sync --extra cu13
uv run python -c "import physicsnemo; print('PhysicsNeMo 版本:', physicsnemo.__version__)"

要安装带有可选功能扩展（例如 nn-extras）的包：

uv sync --extra cu13 --extra nn-extras

对于 CUDA 12 环境，请将 cu13 替换为 cu12：

uv sync --extra cu12 --extra nn-extras

NVCR 容器

可以从 NVIDIA 容器注册表 (NGC) 拉取 PhysicsNeMo Docker 镜像：

NVIDIA 容器注册表（请参考 NGC 注册表以获取最新标签）：

docker pull nvcr.io/nvidia/physicsnemo/physicsnemo:25.06

在容器内，您可以克隆 PhysicsNeMo 的 Git 仓库，并开始使用示例。以下命令展示了如何启动 PhysicsNeMo 容器并从此仓库运行示例：

docker run --shm-size=1g --ulimit memlock=-1 --ulimit stack=67108864 --runtime nvidia \
--rm -it nvcr.io/nvidia/physicsnemo/physicsnemo:25.06 bash
git clone https://github.com/NVIDIA/physicsnemo.git
cd physicsnemo/examples/cfd/darcy_fno/
pip install warp-lang # 安装 NVIDIA Warp 以运行 Darcy 示例
python train_fno_darcy.py

从源代码安装

要从源代码本地构建 PhysicsNeMo Python 包，请执行以下操作：

git clone git@github.com:NVIDIA/physicsnemo.git && cd physicsnemo

pip install --upgrade pip
pip install .
python -c "import physicsnemo; print('PhysicsNeMo 版本:', physicsnemo.__version__)"

对于可编辑安装、本地测试更改以及贡献工作流程，请参阅文档中的 自定义 PhysicsNeMo 指南。此外，请参阅 贡献指南，了解拉取请求、编码规范和 CI 流程，以及 开发者维基，以获取社区和贡献概述。

从源代码构建 Docker 镜像

要构建 PhysicsNeMo Docker 镜像：

docker build -t physicsnemo:deploy \
    --build-arg TARGETPLATFORM=linux/amd64 --target deploy -f Dockerfile .

或者，您也可以运行 make container-deploy。

要构建 CI 镜像：

docker build -t physicsnemo:ci \
    --build-arg TARGETPLATFORM=linux/amd64 --target ci -f Dockerfile .

或者，您也可以运行 make container-ci。

平台支持

对于 pip 或 uv 安装，支持 Linux、macOS（ARM）和 Windows。

Docker 容器仅适用于 linux/amd64 和 linux/arm64 平台。如果使用 linux/arm64，某些依赖项（如 warp-lang）可能无法正确安装。

PhysicsNeMo 迁移指南

NVIDIA Modulus 已更名为 NVIDIA PhysicsNeMo。迁移时请注意：

• 对于 PyPI 轮子，请使用 pip install nvidia-physicsnemo 而不是 pip install nvidia-modulus。
• 对于 Docker 容器，请使用 nvcr.io/nvidia/physicsnemo/physicsnemo:<tag> 而不是 nvcr.io/nvidia/modulus/modulus:<tag>。
• 在您的 pip 依赖文件中（requirements.txt、setup.py、setup.cfg、pyproject.toml 等），将 nvidia-modulus 替换为 nvidia-physicsnemo。
• 在代码中，将导入语句从 import modulus 更改为 import physicsnemo。

旧的 PyPI 注册表和 NGC 容器注册表将很快被弃用，并且不会再接收任何错误修复或更新。不过，旧的检查点仍将与这些更新兼容。

更多详细信息即将发布。

DGL 到 PyTorch Geometric 迁移指南

PhysicsNeMo 支持广泛的图神经网络（GNN），包括 MeshGraphNet 等。目前，PhysicsNeMo 使用 DGL 库作为其 GNN 后端，并计划在未来的版本中完全过渡到 PyTorch Geometric（PyG）。有关详细信息，请参阅 DGL 到 PyG 迁移指南。

为 PhysicsNeMo 做贡献

PhysicsNeMo 是一个开源协作项目，它的成功离不开社区成员对物理与机器学习领域的持续贡献。感谢您为该项目做出贡献，让其他开发者能够在此基础上进一步开发。

有关如何为 PhysicsNeMo 做贡献的指导，请参阅 贡献指南（包括拉取请求、代码风格和持续集成）。如需了解从源码安装、可编辑安装以及测试您的更改等步骤，请参阅文档中的 自定义 PhysicsNeMo 指南。有关社区和贡献的概述，请访问 开发者 Wiki。

引用 PhysicsNeMo

如果您在研究中使用了 PhysicsNeMo，并希望对其进行引用，请参考 引用指南。

沟通交流

• GitHub Discussions：讨论新的架构、实现、物理与机器学习研究等。
• GitHub Issues：用于报告 bug、提出功能请求、解决安装问题等。
• PhysicsNeMo 论坛：PhysicsNeMo 论坛 面向初级至中级用户及开发者，提供通用聊天、在线讨论、协作等功能。

反馈

您想为 PhysicsNeMo 提出改进建议吗？请使用我们的 反馈表单。

许可证

PhysicsNeMo 采用 Apache License 2.0 许可证进行授权。完整的许可证文本请参阅 LICENSE.txt。此外，NVAIE 提供企业级 SLA、技术支持以及预览版访问权限。

NVIDIA PhysicsNeMo 快速上手指南

NVIDIA PhysicsNeMo 是一个开源深度学习框架，专为构建、训练和推理物理人工智能（Physics AI）模型而设计。它结合了先进的科学机器学习（SciML）方法，适用于 AI4Science 和工程领域，支持神经算子、图神经网络（GNN）、Transformer 及物理信息神经网络（PINN）等多种架构。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐 Ubuntu 20.04/22.04)
• GPU: NVIDIA GPU (建议 Ampere 架构或更高版本以获最佳性能)
• CUDA: 已安装与 PyTorch 版本兼容的 CUDA Toolkit
• Python: Python 3.9 - 3.11
• 前置依赖:
  • PyTorch (最新稳定版)
  • NVIDIA Apex (可选，用于混合精度训练加速)
  • Git

注意：PhysicsNeMo 深度集成于 PyTorch 生态，熟悉 PyTorch 的用户将能更快上手。

安装步骤

1. 创建虚拟环境（推荐）

python -m venv physicsnemo-env
source physicsnemo-env/bin/activate

2. 安装 PyTorch

请访问 PyTorch 官网 获取适合您环境的安装命令。例如：

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

(国内用户若遇网络问题，可尝试使用清华源：--index-url https://pypi.tuna.tsinghua.edu.cn/simple)

3. 安装 PhysicsNeMo

目前推荐通过源码安装以获取最新功能（特别是 v2.0 更新）：

git clone https://github.com/NVIDIA/physicsnemo.git
cd physicsnemo
pip install -e .

如果需要安装特定领域的扩展包（如流体力学 CFD 或气象 Earth-2），请参考各子仓库的安装说明。

基本使用

PhysicsNeMo 的核心优势在于其模块化的模型库和分布式训练能力。以下是一个使用预置模型架构进行简单初始化的示例。

示例：初始化一个傅里叶神经算子 (FNO) 模型

此示例展示如何导入并实例化一个经典的 FNO 模型，这是处理偏微分方程求解的常用架构。

import torch
from physicsnemo.models import FNO

# 定义模型参数
in_channels = 3   # 输入通道数
out_channels = 2  # 输出通道数
layers = 4        # 网络层数
hidden_size = 64  # 隐藏层维度

# 实例化模型
model = FNO(
    in_channels=in_channels,
    out_channels=out_channels,
    layers=layers,
    hidden_size=hidden_size
)

# 将模型移至 GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# 准备模拟输入数据 (Batch_Size, Channels, Height, Width)
dummy_input = torch.randn(4, in_channels, 64, 64).to(device)

# 前向推理
with torch.no_grad():
    output = model(dummy_input)

print(f"输入形状：{dummy_input.shape}")
print(f"输出形状：{output.shape}")

下一步

• 数据加载: 使用 physicsnemo.datapipes 加载点云、网格等科学数据。
• 物理约束: 利用 physicsnemo.sym.eq 模块定义偏微分方程（PDEs）作为损失函数的一部分，构建 PINN。
• 分布式训练: 使用 physicsnemo.distributed 模块轻松扩展到多卡或多节点训练。

更多详细教程和案例请参考官方文档：PhysicsNeMo Documentation。