NVIDIA PhysicsNeMo

📝 NVIDIA PhysicsNeMo is undergoing an update to v2.0 - all the features, with easier installation and integration to external packages. See the migration guide for more details!

NVIDIA PhysicsNeMo | Documentation | Install Guide | Getting Started | Contributing Guidelines | Dev blog

What is PhysicsNeMo?

NVIDIA PhysicsNeMo is an open-source deep-learning framework for building, training, fine-tuning, and inferring Physics AI models using state-of-the-art SciML methods for AI4Science and engineering.

PhysicsNeMo provides Python modules to compose scalable and optimized training and inference pipelines to explore, develop, validate, and deploy AI models that combine physics knowledge with data, enabling real-time predictions.

Whether you are exploring the use of neural operators, GNNs, or transformers, or are interested in Physics-Informed Neural Networks or a hybrid approach in between, PhysicsNeMo provides you with an optimized stack that will enable you to train your models at scale.

• More About PhysicsNeMo
  • Scalable GPU-Optimized Training Library
  • A Suite of Physics-Informed ML Models
  • Seamless PyTorch Integration
  • Easy Customization and Extension
  • AI4Science Library
    • Domain-Specific Packages
• Who is Using and Contributing to PhysicsNeMo
• Why Use PhysicsNeMo
• Getting Started
• Resources
• Installation
• Contributing
• Communication
• License

More About PhysicsNeMo

At a granular level, PhysicsNeMo is developed as modular functionality and therefore provides built-in composable modules that are packaged into a few key components:

For a complete list, refer to the PhysicsNeMo API documentation for PhysicsNeMo.

AI4Science Library

Usually, PhysicsNeMo is used either as:

• A complementary tool to PyTorch when exploring AI for SciML and AI4Science applications.
• A deep learning research platform that provides scale and optimal performance on NVIDIA GPUs.

Domain-Specific Packages

The following are packages dedicated to domain experts of specific communities, catering to their unique exploration needs:

• PhysicsNeMo CFD: Inference sub-module of PhysicsNeMo to enable CFD domain experts to explore, experiment, and validate using pretrained AI models for CFD use cases.
• PhysicsNeMo Curator: Inference sub-module of PhysicsNeMo to streamline and accelerate the process of data curation for engineering datasets.
• Earth-2 Studio: Inference sub-module of PhysicsNeMo to enable climate researchers and scientists to explore and experiment with pretrained AI models for weather and climate.

Scalable GPU-Optimized Training Library

PhysicsNeMo provides a highly optimized and scalable training library for maximizing the power of NVIDIA GPUs. Distributed computing utilities allow for efficient scaling from a single GPU to multi-node GPU clusters with a few lines of code, ensuring that large-scale physics-informed machine learning (ML) models can be trained quickly and effectively. The framework includes support for advanced optimization utilities, tailor-made datapipes, and symbolic PDE utilities to enhance end-to-end training speed.

A Suite of Physics-Informed ML Models

PhysicsNeMo offers a library of state-of-the-art models specifically designed for Physics-ML applications. Users can build any model architecture by using the underlying PyTorch layers and combining them with curated PhysicsNeMo layers.

The Model Zoo includes optimized implementations of families of model architectures such as Neural Operators:

• Fourier Neural Operators (FNOs)
• DeepONet
• DoMINO
• Graph Neural Networks (GNNs)
• MeshGraphNet
• MeshGraphNet for Lagrangian
• XAeroNet
• Diffusion Models
• Correction Diffusion Model
• DDPM
• PhysicsNeMo GraphCast
• Transsolver
• RNNs
• SwinVRNN
• Physics-Informed Neural Networks (PINNs)

And many others.

These models are optimized for various physics domains, such as computational fluid dynamics, structural mechanics, and electromagnetics. Users can download, customize, and build upon these models to suit their specific needs, significantly reducing the time required to develop high-fidelity simulations.

Seamless PyTorch Integration

PhysicsNeMo is built on top of PyTorch, providing a familiar and user-friendly experience for those already proficient with PyTorch. This includes a simple Python interface and modular design, making it easy to use PhysicsNeMo with existing PyTorch workflows. Users can leverage the extensive PyTorch ecosystem, including its libraries and tools, while benefiting from PhysicsNeMo's specialized capabilities for physics-ML. This seamless integration ensures users can quickly adopt PhysicsNeMo without a steep learning curve.

For more information, refer to Converting PyTorch Models to PhysicsNeMo Models.

Easy Customization and Extension

PhysicsNeMo is designed to be highly extensible, allowing users to add new functionality with minimal effort. The framework provides Pythonic APIs for defining new physics models, geometries, and constraints, making it easy to extend its capabilities to new use cases. The adaptability of PhysicsNeMo is further enhanced by key features such as ONNX support for flexible model deployment, robust logging utilities for streamlined error handling, and efficient checkpointing to simplify model loading and saving.

This extensibility ensures that PhysicsNeMo can adapt to the evolving needs of researchers and engineers, facilitating the development of innovative solutions in the field of physics-ML.

Detailed information on features and capabilities can be found in the PhysicsNeMo documentation.

Reference samples cover a broad spectrum of physics-constrained and data-driven workflows to suit the diversity of use cases in the science and engineering disciplines.

[!TIP] Have questions about how PhysicsNeMo can assist you? Try our [Experimental] chatbot, PhysicsNeMo Guide, for answers.

Hello World

You can start using PhysicsNeMo in your PyTorch code as simply as shown here:

>>> 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])

To use the distributed module, you can do the following (example for distributed data parallel training; for a more in-depth tutorial, refer to PhysicsNeMo Distributed):

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)

    # Set up the optimizer
    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

    # Sample training loop
    for i in range(20):
        # Random inputs and targets for simplicity
        input = torch.randn(128, 32, device=dist.device)
        target = torch.randn(128, 64, device=dist.device)

        # Training step
        loss = training_step(input, target)

if __name__ == "__main__":
    main()

To use the PDE module, you can do the following:

>>> from physicsnemo.sym.eq.pdes.navier_stokes import NavierStokes
>>> ns = NavierStokes(nu=0.01, rho=1, dim=2)
>>> ns.pprint()
continuity: u__x + v__y
momentum_x: u*u__x + v*u__y + p__x + u__t - 0.01*u__x__x - 0.01*u__y__y
momentum_y: u*v__x + v*v__y + p__y + v__t - 0.01*v__x__x - 0.01*v__y__y

Who is Using and Contributing to PhysicsNeMo

PhysicsNeMo is an open-source project and gets contributions from researchers in the SciML and AI4Science fields. While the PhysicsNeMo team works on optimizing the underlying software stack, the community collaborates and contributes model architectures, datasets, and reference applications so we can innovate in the pursuit of developing generalizable model architectures and algorithms.

Some recent examples of community contributors are the HP Labs 3D Printing team, Stanford Cardiovascular research team, UIUC team, CMU team, etc.

Recent examples of research teams using PhysicsNe