Accelerating AI Physics development with PyTorch and PhysicsNeMo
This tutorial blog is designed for AI researchers and scientific machine learning (SciML) developers who wish to leverage PhysicsNeMo in creating AI Physics models. NVIDIA PhysicsNeMo is built upon PyTorch, the leading standard for AI model development. Organizations can utilize the native PyTorch stack and PhysicsNeMo to go from proof-of-concept development to production-scale training and inference workflows, achieving significant performance and operational efficiencies through enterprise-hardened PhysicsNeMo modules. This blog will demonstrate how AI researchers and SciML developers can progressively integrate PhysicsNeMo modules into their PyTorch stack.
Assuming you are already building SciML models with PyTorch, one of PhysicsNeMo's core advantages is its ability to accelerate your model development cycle. It provides enterprise-hardened building blocks that allow you to compose your training and inference pipelines. By selecting performant and scalable modules, you can build models faster with fewer iterations.
Let's look at the paths different user personas can take to incrementally leverage PhysicsNeMo modules to address the key challenges for their specific needs. Here are three broad user personas:
- Domain - specific application developers where your objective is to develop a custom model that is suited to your own data and customized to the nuances of your specific application
- SciML researchers where your objective is to innovate and discover new model architectures that represent a breakthrough in methodology and application to new problems
- SciML practitioner - any other engineer, scientist or students learning the application of AI for engineering and scientific applications
Domain - specific application developers
Let's take an example of a CAE developer building an AI augmented CFD application for external automotive aerodynamics. The objective is to train a deep learning surrogate model that can predict the drag coefficient for an automotive geometry.
One of the key activities before custom model development is to evaluate the existing pretrained models for its skill. Now instead of writing a model evaluation baseline from scratch, you can import the PhysicsNeMo-CFD python package into your PyTorch environment and evaluate the baseline pretrained models like DoMINO etc with a baseline dataset like DriverML. The pseudo code below outlines the core steps and you can refer to this Jupyter notebook for benchmarking and documentation to run DoMINO NIM as a starting point.
# psuedo code
pip install physicsnemo-cfd
# infer surface data using any other public pre-trained model to generate surface data
python infer_custom_model.py
# or
python infer_domino_nim.py
# benchmark
python generate_surface_benchmarks.py \
./model_output_data/surface_data \
--field-mapping '{"p":"pMean", "wallShearStress":"wallShearStressMean", "pPred":"pMeanPred", "wallShearStressPred":"wallShearStressMeanPred"}'
The next step would be to pick a specific model architecture to train on custom data. PhysicsNeMo makes this very simple - start with one of the curated training recipes by installing any additional required packages (requirements.txt) . Each of the recipes in the training recipe library provides a how-to for a specific application and has guidance on replicating and customizing the recipe to your custom data. This gets you to a first prototype of your AI solution in a matter of hours.
# pseudo code
# each training recipe has additional requiriments
pip install -r requirements.txt
python train.py
There are other key modules that would be of interest but lets view them in the context of our other user persona - the AI researcher.
SciML Researchers
PyTorch serves as an excellent starting point for AI researchers, particularly for those experimenting with new model architectures. Initially, a model might be trained using a novel architecture and a minimal base dataset.
The subsequent phase of innovation involves retraining the model on a larger and higher-resolution dataset. For example, industrial automotive simulations often utilize meshes with 10 to 100 million nodes. Therefore, a SciML researcher developing an architecture for industrial automotive applications must be able to handle such large-scale input meshes. PhysicsNeMo addresses this by offering a distributed manager that employs a "novel domain parallel" strategy to parallelize over extensive inputs.
SciML researchers can achieve scalability by integrating the PhysicsNeMo distributed manager into their PyTorch codebase.Easy to use distributed tooling to initialize and use the torch distributed package in a variety of production environments saves you from the boilerplate of distributed startup, and works across a variety of process launchers and production environments. Furthermore the differentiable distributed functionality, with both differentiable single-layer functions like scatter_v as well as drop-in domain parallel tools like ShardTensor allows you to parallelize over extremely high resolution data, at both training and inference time. This has been described in detail in the documentation here.
# pseudo code
from physicsnemo.distributed import DistributedManager
from physicsnemo.launch.utils import load_checkpoint, save_checkpoint
from physicsnemo.models.model import model
DistributedManager.initialize()
dm = DistributedManager()
...
device = dm.device
You can also leverage the relevant model architecture blocks that are hardened and optimized, which accelerates your productivity by saving time spent developing and troubleshooting from scratch.
One of the core principles of PhysicsNeMo is to enable the community in developing highly generalizable Physics AI models. This is critical to unlock the use of these models in real world enterprise applications. This is because in the absence of generalizable models, users of enterprise solutions have to take on the onus of training bespoke models, generating data for each simulation scenario which is not a practical workflow both on a computational and expertise front. By importing the PhysicsNeMo Sym modules, AI researchers can experiment with the efficacy of integrating governing equations into their training iterations. The example here showcases use of physicsnemo.sym modules to physics-inform your model training. This has been described in detail in the documentation here.
# pseudo code
from physicsnemo.sym.eq.pdes.navier_stokes import NavierStokes
from physicsnemo.sym.eq.phy_informer import PhysicsInformer
ns = NavierStokes(nu=0.1, rho=1.0, dim=2, time=True)
phy_inf = PhysicsInformer(
required_outputs=["continuity", "momentum_x"],
equations=ns,
grad_method="finite_difference"
)
Beyond the model architectures/layers, distributed tooling, and physics-informing capabilities mentioned previously, PhysicsNeMo also offers utilities for other arduous tasks of the model development process like experiment tracking, logging, and model checkpoint/restart.
SciML practitioners
Now for other SciML practitioners, you can follow any of the above patterns to incrementally incorporate PhysicsNeMo into your PyTorch workflow and accelerate your learning and development. A great starting point for learning to use PhysicsNeMo modules in a SciML training loop is the extensive catalog of PhysicsNeMo examples. To get up and running quickly, PhysicsNeMo also offers a container environment that is a perfect launchpad. Here is more information on steps to get started.
These are just a few ways you can use PhysicsNeMo modules to accelerate your AI-Physics model development incrementally. As a developer, you are free to use any or all of these to your liking, and customize to your preferences. The PhysicsNeMo package has benefited from such customizations, and as an open source project welcomes contributions from the SciML community to incorporate novel architectures, loss functions, layers, or other components they want to share with the wider community.
Takeaways
If you are a SciML developer or a Physics AI researcher or an AI practitioner, PhysicsNeMo is a powerful tool in your arsenal to supercharge and extend your PyTorch stack. Instead of building everything from scratch, you can import PhysicsNeMo modules to develop enterprise scale Physics AI solutions with unprecedented speed and simplicity. You can get started easily and step by step using these resources: