PhysicsNeMo Modules#

Basics#

PhysicsNeMo contains its own Module class for constructing neural networks. This class is built on top of PyTorch’s nn.Module and can be used interchangeably within the PyTorch ecosystem. Using PhysicsNeMo modules allows you to leverage several features aimed at improving ease of use, including:

Portable checkpointing via .mdlus files.
A global model registry for discovering and retrieving model classes by name.
Backward compatibility tooling for evolving model classes without breaking existing checkpoints.

In addition, PhysicsNeMo ships a model zoo of optimized architectures that can be used off-the-shelf or composed into larger models. We discuss each of these features in the following sections. For the full programmatic interface, see the API Reference section at the bottom of this page.

Model Zoo#

PhysicsNeMo ships several optimized, customizable and easy-to-use model architectures. These include general-purpose models like Fourier Neural Operators (FNOs), ResNet, and Graph Neural Networks (GNNs) as well as domain-specific models like Deep Learning Weather Prediction (DLWP) and Spherical Fourier Neural Operators (SFNO). Many of these architectures include built-in performance optimizations.

For a list of currently available models, please refer to the models on GitHub.

Below are some simple examples of how to use these models.

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

>>> import torch
>>> from physicsnemo.models.fno.fno import FNO
>>> model = FNO(
        in_channels=4,
        out_channels=3,
        decoder_layers=2,
        decoder_layer_size=32,
        dimension=2,
        latent_channels=32,
        num_fno_layers=2,
        padding=0,
    )
>>> input = torch.randn(32, 4, 32, 32) #(N, C, H, W)
>>> output = model(input)
>>> output.size()
torch.Size([32, 3, 32, 32])

How to Write Your Own PhysicsNeMo Model#

There are a few different ways to construct a PhysicsNeMo model. If you are a seasoned PyTorch user, the easiest way would be to write your model using the optimized layers and utilities from PhysicsNeMo or PyTorch. Let’s take a look at a simple example of a UNet model first showing a simple PyTorch implementation and then the same model rewritten as a PhysicsNeMo module.

import torch.nn as nn

class UNet(nn.Module):
    def __init__(self, in_channels=1, out_channels=1):
        super(UNet, self).__init__()

        self.enc1 = self.conv_block(in_channels, 64)
        self.enc2 = self.conv_block(64, 128)

        self.dec1 = self.upconv_block(128, 64)
        self.final = nn.Conv2d(64, out_channels, kernel_size=1)

    def conv_block(self, in_channels, out_channels):
        return nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2)
        )

    def upconv_block(self, in_channels, out_channels):
        return nn.Sequential(
            nn.ConvTranspose2d(in_channels, out_channels, 2, stride=2),
            nn.Conv2d(out_channels, out_channels, 3, padding=1),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        x1 = self.enc1(x)
        x2 = self.enc2(x1)
        x = self.dec1(x2)
        return self.final(x)

To turn this into a PhysicsNeMo model, the only required change is to inherit from Module instead of torch.nn.Module:

import physicsnemo

class UNet(physicsnemo.Module):  # the only change
    def __init__(self, in_channels=1, out_channels=1):
        super().__init__()
        # ... same code as above ...

Converting PyTorch Models to PhysicsNeMo Models#

In the above example we show constructing a PhysicsNeMo model from scratch. However, you can also convert existing PyTorch models to PhysicsNeMo models in order to leverage PhysicsNeMo features. To do this, use the from_torch() class method as shown below.

import physicsnemo
import torch.nn as nn

class TorchModel(nn.Module):
    def __init__(self):
        super(TorchModel, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = self.conv1(x)
        return self.conv2(x)

# from_torch returns a *class* (not an instance).
# By default, the new class name matches the PyTorch class name ('TorchModel').
PhysicsNeMoModel = physicsnemo.Module.from_torch(TorchModel)
PhysicsNeMoModel.__name__  # 'TorchModel'

# You can override the class name with the ``name`` parameter.
PhysicsNeMoModel = physicsnemo.Module.from_torch(TorchModel, name="MyConvNet")
PhysicsNeMoModel.__name__  # 'MyConvNet'

# Once instantiated, the result is a PhysicsNeMo Module whose class name
# is the one specified above.
model = PhysicsNeMoModel()

Optionally, you can register the converted class in the model registry by passing register=True. This is useful if you later want to load the model from a checkpoint by name, but it is not strictly required — from_checkpoint() can also resolve the class by its module path. You can achieve the same result by calling register() directly.

PhysicsNeMoModel = physicsnemo.Module.from_torch(
    TorchModel, name="MyConvNet", register=True
)

Importing Models from Third-Party Libraries

The same approach works for models defined in third-party libraries such as timm or any other library that provides torch.nn.Module subclasses:

import physicsnemo
import timm

# Load a pre-trained model from timm
TimmResNet = timm.create_model("resnet18", pretrained=False).__class__

# Convert to a PhysicsNeMo Module class
PNMResNet = physicsnemo.Module.from_torch(
    TimmResNet, name="TimmResNet18", register=True
)

# Instantiate as a PhysicsNeMo model
model = PNMResNet()

Saving and Loading PhysicsNeMo Models#

PhysicsNeMo models are interoperable with PyTorch models. You can save and load them using the standard PyTorch APIs, but PhysicsNeMo also provides utilities that save the full constructor arguments alongside the weights into a single .mdlus checkpoint file. This means a .mdlus file is fully self-describing: it contains everything needed to re-create the model without knowing its class or constructor arguments ahead of time.

There are two ways to load from a .mdlus checkpoint:

Into an already instantiated model: use
save() and load(), just like a regular PyTorch nn.Module. This only transfers the weights (state_dict).
Instantiate and load in one step: use
from_checkpoint(). This resolves the class, creates the instance from the saved constructor arguments, and loads the weights. When the class is known ahead of time, calling KnownClass.from_checkpoint(...) is preferred over Module.from_checkpoint(...) because it avoids a class-resolution step.

Example 1: save and load into an existing instance:

 >>> from physicsnemo.models.mlp.fully_connected import FullyConnected
 >>> model = FullyConnected(in_features=32, out_features=64)
 >>> model.save("model.mdlus") # Save model to .mdlus file
 >>> model.load("model.mdlus") # Load model weights from .mdlus file from already instantiated model
 >>> model
 FullyConnected(
  (layers): ModuleList(
    (0): FCLayer(
      (activation_fn): SiLU()
      (linear): Linear(in_features=32, out_features=512, bias=True)
    )
    (1-5): 5 x FCLayer(
      (activation_fn): SiLU()
      (linear): Linear(in_features=512, out_features=512, bias=True)
    )
  )
  (final_layer): FCLayer(
    (activation_fn): Identity()
    (linear): Linear(in_features=512, out_features=64, bias=True)
  )
)

Example 2: instantiate directly from a checkpoint:

In this case we don’t need to know the class or the constructor arguments. The .mdlus file contains all the information needed to instantiate the model.

 >>> from physicsnemo import Module
 >>> fc_model = Module.from_checkpoint("model.mdlus") # Instantiate model from .mdlus file.
 >>> fc_model
 FullyConnected(
  (layers): ModuleList(
    (0): FCLayer(
      (activation_fn): SiLU()
      (linear): Linear(in_features=32, out_features=512, bias=True)
    )
    (1-5): 5 x FCLayer(
      (activation_fn): SiLU()
      (linear): Linear(in_features=512, out_features=512, bias=True)
    )
  )
  (final_layer): FCLayer(
    (activation_fn): Identity()
    (linear): Linear(in_features=512, out_features=64, bias=True)
  )
)

Note

In order to use from_checkpoint(), the model must have .json-serializable inputs to the __init__ function. The only exception is when the argument is itself a Module instance. In this case, it is possible to construct, save and load nested Modules with multiple levels of nesting and/or multiple Module instances at each level. See Constructing Nested Modules for details. It is highly recommended that all PhysicsNeMo models be developed with this requirement in mind.
Using Module.from_checkpoint will not work if the model has any buffers or parameters that are registered outside of the model’s __init__ function. In that case, use load() instead, or ensure that all model parameters and buffers are registered inside __init__.

For training-loop checkpointing that also saves optimizer, scheduler, and scaler state, see save_checkpoint() and load_checkpoint().

Constructing Nested Modules#

PhysicsNeMo supports constructing nested modules where one Module can accept another Module as an argument to its __init__ function. This allows you to build complex, modular architectures while still benefiting from PhysicsNeMo’s checkpointing and model management features.

Simple Nesting with PhysicsNeMo Modules

The simplest case is nesting Module instances directly:

import torch
import physicsnemo

class EncoderModule(physicsnemo.Module):
    def __init__(self, input_size, hidden_size):
        super().__init__()
        self.encoder = torch.nn.Linear(input_size, hidden_size)
        self.input_size = input_size
        self.hidden_size = hidden_size

    def forward(self, x):
        return self.encoder(x)

class DecoderModule(physicsnemo.Module):
    def __init__(self, hidden_size, output_size):
        super().__init__()
        self.decoder = torch.nn.Linear(hidden_size, output_size)
        self.hidden_size = hidden_size
        self.output_size = output_size

    def forward(self, x):
        return self.decoder(x)

class AutoEncoder(physicsnemo.Module):
    def __init__(self, encoder, decoder):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder

    def forward(self, x):
        encoded = self.encoder(x)
        return self.decoder(encoded)

# Create nested model
encoder = EncoderModule(input_size=64, hidden_size=32)
decoder = DecoderModule(hidden_size=32, output_size=64)
model = AutoEncoder(encoder=encoder, decoder=decoder)

# Save and load with full structure preserved
model.save("autoencoder.mdlus")
loaded_model = physicsnemo.Module.from_checkpoint("autoencoder.mdlus")

Nesting Converted PyTorch Modules

You can also nest PyTorch nn.Module instances, but they must first be converted to Module using from_torch(). All nested PyTorch modules must be converted:

import torch.nn as nn
import physicsnemo

# Define PyTorch modules
class TorchEncoder(nn.Module):
    def __init__(self, input_size, hidden_size):
        super().__init__()
        self.encoder = nn.Linear(input_size, hidden_size)
        self.input_size = input_size
        self.hidden_size = hidden_size

    def forward(self, x):
        return self.encoder(x)

class TorchDecoder(nn.Module):
    def __init__(self, hidden_size, output_size):
        super().__init__()
        self.decoder = nn.Linear(hidden_size, output_size)
        self.hidden_size = hidden_size
        self.output_size = output_size

    def forward(self, x):
        return self.decoder(x)

# Convert to PhysicsNeMo modules
PNMEncoder = physicsnemo.Module.from_torch(TorchEncoder)
PNMDecoder = physicsnemo.Module.from_torch(TorchDecoder)

# Define top-level model
class AutoEncoder(physicsnemo.Module):
    def __init__(self, encoder, decoder):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder

    def forward(self, x):
        encoded = self.encoder(x)
        return self.decoder(encoded)

# Create nested model with converted modules
encoder = PNMEncoder(input_size=64, hidden_size=32)
decoder = PNMDecoder(hidden_size=32, output_size=64)
model = AutoEncoder(encoder=encoder, decoder=decoder)

# Save and load
model.save("autoencoder.mdlus")
loaded_model = physicsnemo.Module.from_checkpoint("autoencoder.mdlus")

What Does NOT Work

You cannot directly pass a torch.nn.Module instance to a Module’s __init__ without converting it first:

# This will NOT work and raise an error during save/load:
class AutoEncoder(physicsnemo.Module):
    def __init__(self, encoder):
        super().__init__()
        self.encoder = encoder  # encoder is a torch.nn.Module

torch_encoder = TorchEncoder(input_size=64, hidden_size=32)
model = AutoEncoder(encoder=torch_encoder)  # This creates the model

# But this will fail:
model.save("autoencoder.mdlus")
# Error: Cannot serialize torch.nn.Module arguments.
# You must use Module.from_torch() to convert it first.

Backward Compatibility#

When evolving a model class over time, for example, renaming the class, adding or removing constructor arguments, you may still need to load checkpoints that were saved with an older version of the class. PhysicsNeMo provides a versioning and argument-mapping system for this purpose.

The key ingredients are:

``__model_checkpoint_version__``: A version string (e.g. "0.2.0") on each Module subclass. This is saved into every .mdlus file and compared at load time.
``__supported_model_checkpoint_version__``: A dict mapping older version strings to warning messages. If the checkpoint version is in this dict, loading proceeds (with a warning) instead of raising an error.
``_backward_compat_arg_mapper``: A classmethod that transforms constructor arguments from an older version into the format expected by the current version.
``_overridable_args`` and ``override_args``: Allow callers of from_checkpoint() to override specific constructor arguments at load time. Only arguments listed in _overridable_args can be overridden.

Example workflow:

Suppose you have an initial model that has been deployed and trained checkpoints exist:

import physicsnemo

class MyModel(physicsnemo.Module):
    __model_checkpoint_version__ = "0.1.0"

    def __init__(self, img_channels, hidden_dim=64):
        super().__init__()
        self.img_channels = img_channels
        self.hidden_dim = hidden_dim
        # ... model layers ...

    def forward(self, x):
        return x

# A trained checkpoint is saved
model = MyModel(img_channels=3)
model.save("my_model.mdlus")

Later, you refactor the class: you rename img_channels to in_channels and remove the hidden_dim argument. To still be able to load old checkpoints:

class MyModel(physicsnemo.Module):
    __model_checkpoint_version__ = "0.2.0"
    __supported_model_checkpoint_version__ = {
        "0.1.0": (
            "Loading MyModel checkpoint from version 0.1.0. "
            "Consider re-saving to upgrade to 0.2.0."
        ),
    }

    @classmethod
    def _backward_compat_arg_mapper(cls, version, args):
        args = super()._backward_compat_arg_mapper(version, args)
        if version == "0.1.0":
            # Rename img_channels -> in_channels
            if "img_channels" in args:
                args["in_channels"] = args.pop("img_channels")
            # Remove deprecated argument
            args.pop("hidden_dim", None)
        return args

    def __init__(self, in_channels):
        super().__init__()
        self.in_channels = in_channels
        # ... updated model layers ...

    def forward(self, x):
        return x

# Old checkpoint loads successfully (with a warning)
loaded = MyModel.from_checkpoint("my_model.mdlus")

For the full details of from_checkpoint(), refer to the API Reference.

PhysicsNeMo Model Registry and Entry Points#

PhysicsNeMo contains a ModelRegistry that provides a single, global lookup for all model classes in the PhysicsNeMo ecosystem. This is useful because:

Stable access: Once a class is registered, it can be retrieved by name regardless of where its source module lives. If the class is later moved to a different package or module path, the registry-based import remains the same.
Checkpoint loading: from_checkpoint() uses the registry to resolve class names stored in .mdlus files.
Third-party integration: External packages can expose models via Python entry points so that they appear in the registry when installed.

A few important rules:

Names must be unique. Attempting to register a name that is already in use raises a ValueError.
The class name is used by default. When registering via register() or Module.from_torch(..., register=True), the __name__ attribute of the class is used unless an explicit name is provided.

Below is a simple example of how to use the model registry to obtain a model class.

>>> from physicsnemo.registry import ModelRegistry
>>> model_registry = ModelRegistry()
>>> model_registry.list_models()
['AFNO', 'DLWP', 'FNO', 'FullyConnected', 'GraphCastNet', 'MeshGraphNet', 'One2ManyRNN', 'Pix2Pix', 'SFNO', 'SRResNet']
>>> FullyConnected = model_registry.factory("FullyConnected")
>>> model = FullyConnected(in_features=32, out_features=64)

You can also register classes automatically when defining them by passing register=True as a class keyword argument:

class MyModel(physicsnemo.Module, register=True):
    def __init__(self, hidden_dim=64):
        super().__init__()
        self.hidden_dim = hidden_dim

    def forward(self, x):
        return x

# The class is now accessible from the registry
registry = ModelRegistry()
ModelClass = registry.factory("MyModel")

Exposing Models via Entry Points

The model registry also allows exposing models via Python entry points. This allows for integration of models into the PhysicsNeMo ecosystem from external packages. For example, suppose you have a package MyPackage that contains a model MyModel. You can expose this model to the PhysicsNeMo registry by adding an entry point to your toml file:

# setup.py

from setuptools import setup, find_packages

setup()

# pyproject.toml

[build-system]
requires = ["setuptools", "wheel"]
build-backend = "setuptools.build_meta"

[project]
name = "MyPackage"
description = "My Neural Network Zoo."
version = "0.1.0"

[project.entry-points."physicsnemo.models"]
MyPhysicsNeMoModel = "mypackage.models:MyPhysicsNeMoModel"

# mypackage/models.py

import torch.nn as nn
from physicsnemo.core import Module

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = self.conv1(x)
        return self.conv2(x)

MyPhysicsNeMoModel = Module.from_torch(MyModel)

Once this package is installed, you can access the model via the PhysicsNeMo model registry.

>>> from physicsnemo.registry import ModelRegistry
>>> model_registry = ModelRegistry()
>>> model_registry.list_models()
['MyPhysicsNeMoModel', 'AFNO', 'DLWP', 'FNO', 'FullyConnected', 'GraphCastNet', 'MeshGraphNet', 'One2ManyRNN', 'Pix2Pix', 'SFNO', 'SRResNet']
>>> MyPhysicsNeMoModel = model_registry.factory("MyPhysicsNeMoModel")

API Reference#

`Module`#

class physicsnemo.core.module.Module(*args, **kwargs)[source]#

The base class for all network models in PhysicsNeMo.

This should be used as a direct replacement for torch.nn.module and provides additional functionality for saving and loading models, as well as handling file system abstractions.

There is one important requirement for all models in PhysicsNeMo. They must have json serializable arguments in their __init__ function. This is required for saving and loading models and allow models to be instantiated from a checkpoint. The only one exception to this rule is when the argument passed to the __init__ function is itself a physicsnemo.Module instance. In this case, it is possible to construct, save and load nested Modules, with multiple levels of nesting and/or multiple physicsnemo.Module instances at each level. To be able to pass a torch.nn.Module instance as an argument to the __init__ function, it is necessary to first use the Module.from_torch method to convert the torch.nn.Module subclass to a physicsnemo.Module subclass To pass nested torch.nn.Module instances as arguments to the __init__ function, it is necessary to convert all nested torch.nn.Module instances to physicsnemo.Module instances using the Module.from_torch method. See the examples below for more details.

Parameters:: meta (ModelMetaData, optional) – Meta data class for storing info regarding model, by default None

Examples

To construct nested physicsnemo.Module instances with multiple levels of nesting and/or multiple physicsnemo.Module instances at each level:

class InnerModel(physicsnemo.Module):
    def __init__(self, hidden_size):
        super().__init__(meta=ModelMetaData())
        self.hidden_size = hidden_size

class OuterModel(physicsnemo.Module):
    def __init__(self, inner_model):
        super().__init__(meta=ModelMetaData())
        self.inner_model = inner_model

# Create and save nested model
model = OuterModel(inner_model=InnerModel(128))
model.save("checkpoint.mdlus")
loaded = physicsnemo.Module.from_checkpoint("checkpoint.mdlus")

Applying this to a torch.nn.Module instance is also possible, as long as all nested torch.nn.Module instances are converted to physicsnemo.Module instances using the Module.from_torch method:

class TorchInnerModel(torch.nn.Module):
    def __init__(self, size):
        super().__init__()
        self.size = size

class TorchMyModel(torch.nn.Module):
    def __init__(self, inner_model):
        super().__init__()
        self.inner_model = inner_model

# Convert both torch.nn.Module to physicsnemo.Module
PNMInnerModel = physicsnemo.Module.from_torch(
    TorchInnerModel, meta=ModelMetaData()
)
PNMMyModel = physicsnemo.Module.from_torch(
    TorchMyModel, meta=ModelMetaData()
)

# Create nested model with converted torch modules
model = PNMMyModel(inner_model=PNMInnerModel(size=128))

When subclassing Module, you can pass register=True as a class argument to automatically register the class in the model registry. This allows the class to be retrieved later by name using the ModelRegistry:

>>> from physicsnemo.core import Module, ModelMetaData, ModelRegistry
>>> class MyCustomModelA(Module, register=True):
...     def __init__(self, hidden_dim=64):
...         super().__init__(meta=ModelMetaData())
...         self.hidden_dim = hidden_dim
>>> # The class is now registered and can be retrieved by name
>>> registry = ModelRegistry()
>>> ModelClass = registry.factory('MyCustomModelA')
>>> model = ModelClass(hidden_dim=128)
>>> model.hidden_dim
128

property device: device#

Get device model is on

Returns:: PyTorch device
Return type:: torch.device

classmethod from_checkpoint( file_name: str, override_args: Dict[str, Any] | None = None, strict: bool = True, ) → Module[source]#

Utility class method for instantiating and loading a Module instance from a ‘.mdlus’ checkpoint file.

Parameters:

file_name (str) – Checkpoint file name. Must be a valid ‘.mdlus’ checkpoint file.
override_args (Optional[Dict[str, Any]], optional, default=None) –
Dictionary of arguments to override the __init__ method’s arguments saved in the checkpoint. The override of arguments occurs before the model is instantiated, which allows for ad-hoc modifications to the model’s initialization. Argument overrides are however applied before the state-dict is loaded, which means that for parameters or buffers saved in the state-dict, the values contained in the state-dict will take precedence over the override. This might also result in unexpected behavior if the model is instantiated with different arguments than the ones saved in the checkpoint, and some mismatching keys are saved in the state-dict.

Note: Only arguments defined in cls._overridable_args can be overridden. Module’s subclasses by default disable this functionality, unless they explicity define an _overridable_args class attribute. Attempting to override any other argument will raise a ValueError. This API should be used with caution and only if you fully understand the implications of the override.
strict (bool, optional) – Whether to strictly enforce that the keys in state_dict match, by default True

Return type:

Module

Raises:

IOError – If file_name provided does not exist or is not a valid checkpoint

Examples

Simple argument override:

class MyModel(Module):
    _overridable_args = set(["a", "b"])
    def __init__(self, a, b=2.0):
        super().__init__()
        # ... model implementation ...
model = MyModel(1.0, b=2.0)
model.save("checkpoint.mdlus")
model_loaded = MyModel.from_checkpoint("checkpoint.mdlus", override_args={"a": 5.0})

For nested module, override is possible with dot-separated syntax:

class SubModule(Module):
    _overridable_args = set(["a"])
    def __init__(self, a):
        super().__init__()
        # ... submodule implementation ...
class MyModel(Module):
    def __init__(self, submodule):
        super().__init__()
        self.submodule = submodule
        # ... model implementation ...
submodule = SubModule(1.0)
model = MyModel(submodule)
model.save("checkpoint.mdlus")
model = MyModel.from_checkpoint("checkpoint.mdlus", override_args={"submodule.a": 2.0})

static from_torch( torch_model_class: type[Module], meta: ModelMetaData | None = None, name: str | None = None, register: bool = False, ) → type[Module][source]#

Construct a PhysicsNeMo module from a PyTorch module. The resulting class is a PhysicsNeMo Module class. Any instance of this class will be a PhysicsNeMo Module instance with an attribute inner_model that is an instance of the PyTorch model class.

Parameters:

torch_model_class (torch.nn.Module) – PyTorch module class
meta (ModelMetaData, optional, default=None) – Meta data for the model.
name (str, optional, default=None) – Name of the PhysicsNeMo model class. Used for registering the class in the model registry. If None, the name of the PyTorch model class is used.
register (bool, optional, default=False) –
Whether to register the class in the model registry. If True, the class will be registered and can be retrieved later using ModelRegistry().factory(name).

Important

To be able to later load the model with from_checkpoint(), it is necessary to register the class in the model registry by setting register=True. A class created via from_torch that is not registered will not be able to be loaded with from_checkpoint.

Return type:

Module

Examples

Example 1: Convert a PyTorch model to PhysicsNeMo without specifying a name:

>>> import torch
>>> import torch.nn as nn
>>> from physicsnemo.core import Module, ModelMetaData, ModelRegistry
>>> # Define a simple MLP in PyTorch
>>> class SimpleMLP(nn.Module):
...     def __init__(self, input_size, hidden_size, output_size):
...         super().__init__()
...         self.input_size = input_size
...         self.hidden_size = hidden_size
...         self.output_size = output_size
...         self.fc1 = nn.Linear(input_size, hidden_size)
...         self.relu = nn.ReLU()
...         self.fc2 = nn.Linear(hidden_size, output_size)
...
...     def forward(self, x):
...         x = self.fc1(x)
...         x = self.relu(x)
...         x = self.fc2(x)
...         return x
>>> # Convert PyTorch model to PhysicsNeMo Module without registering
>>> PNMSimpleMLP = Module.from_torch(SimpleMLP, meta=ModelMetaData())
>>> # The physicsnemo class name is the same as the PyTorch class name
>>> PNMSimpleMLP.__name__
'SimpleMLP'
>>> # Instantiate the PhysicsNeMo model
>>> model = PNMSimpleMLP(input_size=10, hidden_size=64, output_size=5)
>>> # Access the inner PyTorch model
>>> assert model.inner_model.input_size == 10
>>> assert model.inner_model.hidden_size == 64
>>> assert model.inner_model.output_size == 5
>>> # Use the model for inference
>>> x = torch.randn(32, 10)
>>> output = model(x)
>>> output.shape
torch.Size([32, 5])
>>> # Cannot retrieve the model class from the registry because it is not registered

Example 2: Convert a PyTorch model with a custom name:

>>> import torch
>>> import torch.nn as nn
>>> from physicsnemo.core import Module, ModelMetaData, ModelRegistry
>>> # Define a simple MLP in PyTorch
>>> class SimpleMLP(nn.Module):
...     def __init__(self, input_size, hidden_size, output_size):
...         super().__init__()
...         self.input_size = input_size
...         self.hidden_size = hidden_size
...         self.output_size = output_size
...         self.fc1 = nn.Linear(input_size, hidden_size)
...         self.relu = nn.ReLU()
...         self.fc2 = nn.Linear(hidden_size, output_size)
...
...     def forward(self, x):
...         x = self.fc1(x)
...         x = self.relu(x)
...         x = self.fc2(x)
...         return x
>>> # Convert with a custom name
>>> PNMSimpleMLP = Module.from_torch(
...     SimpleMLP,
...     meta=ModelMetaData(),
...     name='CustomSimpleMLP'
... )
>>> # The physicsnemo class name is defined by the name parameter
>>> PNMSimpleMLP.__name__
'CustomSimpleMLP'
>>> # Instantiate the PhysicsNeMo model
>>> model = PNMSimpleMLP(input_size=10, hidden_size=64, output_size=5)
>>> # Access the inner PyTorch model
>>> assert model.inner_model.input_size == 10
>>> assert model.inner_model.hidden_size == 64
>>> assert model.inner_model.output_size == 5
>>> # Cannot retrieve the model class from the registry because it is not registered

>>> import torch
>>> import torch.nn as nn
>>> from physicsnemo.core import Module, ModelMetaData, ModelRegistry
>>> # Define a simple MLP in PyTorch
>>> class SimpleMLP(nn.Module):
...     def __init__(self, input_size, hidden_size, output_size):
...         super().__init__()
...         self.input_size = input_size
...         self.hidden_size = hidden_size
...         self.output_size = output_size
...         self.fc1 = nn.Linear(input_size, hidden_size)
...         self.relu = nn.ReLU()
...         self.fc2 = nn.Linear(hidden_size, output_size)
...
...     def forward(self, x):
...         x = self.fc1(x)
...         x = self.relu(x)
...         x = self.fc2(x)
...         return x
>>> # Convert with register=True to add to the model registry
>>> PNMSimpleMLP = Module.from_torch(
...     SimpleMLP,
...     meta=ModelMetaData(),
...     name='RegisteredMLP',
...     register=True
... )
>>> # The class is now registered and can be retrieved by name
>>> registry = ModelRegistry()
>>> ModelClass = registry.factory('RegisteredMLP')
>>> isinstance(ModelClass, type) and issubclass(ModelClass, Module)
True
>>> # The class name is defined by the name parameter
>>> ModelClass.__name__
'RegisteredMLP'
>>> # Instantiate the model from the registry
>>> model = ModelClass(input_size=10, hidden_size=64, output_size=5)
>>> model.inner_model.input_size
10

classmethod instantiate( arg_dict: Dict[str, Any], ) → Module[source]#

Instantiate a model from a dictionary of arguments. This method is reserved for advanced and internal use cases. For most use cases, it is recommended to instantiate the model using standard instantiation mechanisms.

Parameters:: arg_dict (Dict[str, Any]) – Dictionary of arguments to instantiate model with. This should be have three keys: ‘__name__’, ‘__module__’, and ‘__args__’. The first two are used to import the class and the last is used to instantiate the class. The ‘__args__’ key should be a dictionary of arguments to pass to the class’s __init__ function.
Return type:: Module

Examples

>>> import warnings
>>> warnings.filterwarnings("ignore")
>>> from physicsnemo.core.module import Module
>>> # Define the argument dictionary with the three required keys
>>> arg_dict = {
...     '__name__': 'FullyConnected',
...     '__module__': 'physicsnemo.models.mlp.fully_connected',
...     '__args__': {'in_features': 10, 'out_features': 5}
... }
>>> # Instantiate the model using the class method
>>> model = Module.instantiate(arg_dict)
>>> # Verify the model was created with the correct parameters
>>> x = torch.randn(100, 10)
>>> output = model(x)
>>> output.shape
torch.Size([100, 5])

load( file_name: str, map_location: None | str | device = None, strict: bool = True, ) → None[source]#

Utility method for loading the model weights from a ‘.mdlus’ checkpoint file. Unlike from_checkpoint(), this method does not instantiate the model, but rather loads the state_dict for an already instantiated model.

Parameters:

file_name (str) – Checkpoint file name. Must be a valid ‘.mdlus’ checkpoint file.
map_location (Union[None, str, torch.device], optional, default=None) – Map location for loading the model weights, None will use the model’s device.
strict (bool, optional, default=True) – Whether to strictly enforce that the keys in state_dict match.

Raises:

IOError – If file_name provided does not exist or is not a valid checkpoint

Examples

Basic example loading the model weights (state_dict) from a checkpoint:

from physicsnemo.models.mlp import FullyConnected

# Create a model with the same architecture as the saved one
model = FullyConnected(in_features=32, out_features=64)

# Load the weights from checkpoint
model.load("FullyConnected.mdlus")

Loading with specific device mapping:

import torch
from physicsnemo.models.mlp import FullyConnected

model = FullyConnected(in_features=32, out_features=64)

# Load checkpoint to CPU even if it was saved on GPU
model.load("FullyConnected.mdlus", map_location="cpu")

# Or load to a specific GPU
model.load("FullyConnected.mdlus", map_location=torch.device("cuda:0"))

num_parameters() → int[source]#: Gets the number of learnable parameters

save( file_name: str | None = None, verbose: bool = False, legacy_format: bool = False, ) → None[source]#

Utility method for saving a Module instance to a ‘.mdlus’ checkpoint file.

Parameters:

file_name (Union[str,None], optional, default=None) – File name to save the model checkpoint to. When None is provided it will default to the model’s class name.
verbose (bool, optional, default=False) – Whether to save the model in verbose mode which will include git hash, etc.
legacy_format (bool, optional, default=False) – Whether to save the model in legacy tar format. If True, saves as tar archive. If False (default), saves as zip archive.

Raises:

ValueError – If file_name does not end with .mdlus extension

Examples

>>> from physicsnemo.models.mlp import FullyConnected
>>> model = FullyConnected(in_features=32, out_features=64)
>>> # Save a checkpoint with the default file name 'FullyConnected.mdlus' (using the class name).
>>> model.save()
>>> # Save a checkpoint to a specified file name 'my_model.mdlus'
>>> model.save("my_model.mdlus")

`ModelRegistry`#

class physicsnemo.core.registry.ModelRegistry(*args, **kwargs)[source]#

factory(name: str) → type['Module'][source]#

Returns a registered model class given its name.

Parameters:: name (str) – The name of the registered model.
Returns:: model – The registered model.
Return type:: physicsnemo.core.Module
Raises:: KeyError – If no model is registered under the provided name.

list_models() → List[str][source]#

Returns a list of the names of all models currently registered in the registry.

Returns:: A list of the names of all registered models. The order of the names is not guaranteed to be consistent.
Return type:: List[str]

register( model: type['Module'], name: str | None = None, ) → None[source]#

Registers a physicsnemo model class in the model registry under the provided name. If no name is provided, the model’s name (from its __name__ attribute) is used. If the name is already in use, raises a ValueError.

Parameters:

model (physicsnemo.core.Module) – The model class to be registered.
name (str, optional) – The name to register the model under. If None, the model class name is used.

Raises:

ValueError – If the provided name is already in use in the registry.

Examples

Example 1: Register a model class using its default name (from __name__):

>>> from physicsnemo.core import Module, ModelRegistry
>>> # Define a custom model class
>>> class MyCustomModel(Module):
...     def __init__(self, hidden_size):
...         super().__init__()
...         self.hidden_size = hidden_size
...
...     def forward(self, x):
...         return x
>>> # Get the registry instance
>>> registry = ModelRegistry()
>>> # Register the model without specifying a name
>>> # The class name 'MyCustomModel' will be used automatically
>>> registry.register(MyCustomModel)
>>> # Retrieve the model class from the registry
>>> ModelClass = registry.factory('MyCustomModel')
>>> # Instantiate the model
>>> model = ModelClass(hidden_size=128)

PhysicsNeMo Modules#

Basics#

Model Zoo#

How to Write Your Own PhysicsNeMo Model#

Converting PyTorch Models to PhysicsNeMo Models#

Saving and Loading PhysicsNeMo Models#

Constructing Nested Modules#

Backward Compatibility#

PhysicsNeMo Model Registry and Entry Points#

API Reference#

Module#

ModelRegistry#

`Module`#

`ModelRegistry`#