Modulus Models - NVIDIA Docs

NVIDIA Docs Hub NVIDIA Modulus NVIDIA Modulus Core (Latest Release) Modulus Models

Basics

Modulus contains its own Model class for constructing neural networks. This model class is built on top of PyTorch’s nn.Module and can be used interchangeably within the PyTorch ecosystem. Using Modulus models allows you to leverage various features of Modulus aimed at improving performance and ease of use. These features include, but are not limited to, model zoo, automatic mixed-precision, CUDA Graphs, and easy checkpointing. We discuss each of these features in the following sections.

Model Zoo

Modulus contains several optimized, customizable and easy-to-use models. These include some very general 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).

Currently available models include:

Model Name	Inputs	Outputs
FullyConnected	torch.Tensor [N, in_features]	torch.Tensor [N, out_features]
FourierNeuralOperator	torch.Tensor [N, in_channels, H, W]	torch.Tensor [N, out_channels, H, W]
AdaptiveFourierNeuralOperator	torch.Tensor [N, in_channels, H, W]	torch.Tensor [N, out_channels, H, W]
MeshGraphNet	torch.Tensor [num_nodes, input_dim_nodes], torch.Tensor [num_edges, input_dim_edges], dgl.DGLGraph [num_nodes, num_edges]	torch.Tensor [num_nodes, output_dim]
GraphCastNet	torch.Tensor [N, C_in, H, W]	torch.Tensor [N, C_out, H, W]
Pix2PixNet	torch.Tensor [N, in_channels, H, W]	torch.Tensor [N, out_channels, H, W]
One2ManyRNN	torch.Tensor [N, C, 1, H, W]	torch.Tensor [N, C, T, H, W]
Seq2SeqRNN	torch.Tensor [N, C, T, H, W]	torch.Tensor [N, C, T, H, W]
SRResNet	torch.Tensor [N, C_in, D, H, W]	torch.Tensor [N, C_out, D_out, H_out, W_out]
DLWP	torch.Tensor [N, C_in, 6, Res, Res]	torch.Tensor [N, C_out, 6, Res, Res]
SphericalFourierNeuralOperatorNet	torch.Tensor [N, C_in, H, W]	torch.Tensor [N, C_out, H, W]

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

Copy
Copied!

            
            >>> import torch
>>> from modulus.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])

Copy
Copied!

            
            >>> import torch
>>> from modulus.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 Modulus model

There are a few different ways to construct a Modulus model. If you are a seasoned PyTorch user, the easiest way would be to write your model using the optimized layers and utilities from Modulus or Pytorch. Lets take a look at a simple example of a UNet model first showing a simple PyTorch implementation and then a Modulus implementation that supports CUDA Graphs and Automatic Mixed-Precision.

Copy
Copied!

            
            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)

Now we show this model rewritten in Modulus. First, let’s subclass the model from modulus.Module instead of torch.nn.Module. The modulus.Module class acts like a direct replacement for the torch.nn.Module and provides additional functionality for saving and loading checkpoints, etc. Refer to the API docs of modulus.Module for further details. Additionally we will add metadata to the model to capture the optimizations that this model supports. In this case we will enable CUDA Graphs and Automatic Mixed-Precision.

Copy
Copied!

            
            from dataclasses import dataclass
import modulus
import torch.nn as nn

@dataclass
class UNetMetaData(modulus.ModelMetaData):
    name: str = "UNet"
    # Optimization
    jit: bool = True
    cuda_graphs: bool = True
    amp_cpu: bool = True
    amp_gpu: bool = True

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

        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)

Now that we have our Modulus model, we can make use of these optimizations using the modulus.utils.StaticCaptureTraining decorator. This decorator will capture the training step function and optimize it for the specified optimizations.

Copy
Copied!

            
            import torch
from modulus.utils import StaticCaptureTraining

model = UNet().to("cuda")
input = torch.randn(8, 1, 128, 128).to("cuda")
output = torch.zeros(8, 1, 64, 64).to("cuda")

optim = torch.optim.Adam(model.parameters(), lr=0.001)

# Create training step function with optimization wrapper
# StaticCaptureTraining calls `backward` on the loss and
# `optimizer.step()` so you don't have to do that
# explicitly.
@StaticCaptureTraining(
    model=model,
    optim=optim,
    cuda_graph_warmup=11,
)
def training_step(invar, outvar):
    predvar = model(invar)
    loss = torch.sum(torch.pow(predvar - outvar, 2))
    return loss

# Sample training loop
for i in range(20):
    # In place copy of input and output to support cuda graphs
    input.copy_(torch.randn(8, 1, 128, 128).to("cuda"))
    output.copy_(torch.zeros(8, 1, 64, 64).to("cuda"))

    # Run training step
    loss = training_step(input, output)

For the simple model above, you can observe ~1.1x speed-up due to CUDA Graphs and AMP. The speed-up observed changes from model to model and is typically greater for more complex models.

Note

The ModelMetaData and modulus.Module do not make the model support CUDA Graphs, AMP, etc. optimizations automatically. The user is responsible to write the model code that enables each of these optimizations. Models in the Modulus Model Zoo are written to support many of these optimizations and checked against Modulus’s CI to ensure that they work correctly.

Note

The StaticCaptureTraining decorator is still under development and may be refactored in the future.

Converting PyTorch Models to Modulus Models

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

Copy
Copied!

            
            from dataclasses import dataclass
import modulus
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)

@dataclass
class ConvMetaData(ModelMetaData):
    name: str = "UNet"
    # Optimization
    jit: bool = True
    cuda_graphs: bool = True
    amp_cpu: bool = True
    amp_gpu: bool = True

ModulusModel = modulus.Module.from_torch(TorchModel, meta=ConvMetaData())

Saving and Loading Modulus Models

As mentioned above, Modulus models are interoperable with PyTorch models. This means that you can save and load Modulus models using the standard PyTorch APIs however, we provide a few additional utilities to make this process easier. A key challenge in saving and loading models is keeping track of the model metadata such as layer sizes, etc. Modulus models can be saved with this metadata to a custom .mdlus file. These files allow for easy loading and instantiation of the model. We show two examples of this below. The first example shows saving and loading a model from an already instantiated model.

Copy
Copied!

            
             >>> from modulus.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)
  )
)

The second example shows loading a model from a .mdlus file without having to instantiate the model first. We note that in this case we don’t know the class or parameters to pass to the constructor of the model. However, we can still load the model from the .mdlus file.

Copy
Copied!

            
             >>> from modulus 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 make use of this functionality, the model must have json serializable inputs to the __init__ function. It is highly recommended that all Modulus models be developed with this requirement in mind.

Modulus Model Registry and Entry Points

Modulus contains a model registry that allows for easy access and ingestion of models. Below is a simple example of how to use the model registry to obtain a model class.

Copy
Copied!

            
            >>> from modulus.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)

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

Copy
Copied!

            
            # setup.py

from setuptools import setup, find_packages

setup()

Copy
Copied!

            
            # 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."modulus.models"]
MyModulusModel = "mypackage.models.MyModulusModel:MyModulusModel"

Copy
Copied!

            
            # mypackage/models.py

import torch.nn as nn
from modulus.models import Model

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)

MyModulusModel = Model.from_pytorch(MyModel)

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

Copy
Copied!

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

For more information on entry points and potential use cases, see this blog post.

Fully Connected Network

class modulus.models.mlp.fully_connected.FullyConnected(*args, **kwargs)[source]

Bases: Module

A densely-connected MLP architecture

Parameters

Example

Copy
Copied!

            
            >>> model = modulus.models.mlp.FullyConnected(in_features=32, out_features=64)
>>> input = torch.randn(128, 32)
>>> output = model(input)
>>> output.size()
torch.Size([128, 64])

forward(x: Tensor) → Tensor[source]

class modulus.models.mlp.fully_connected.MetaData(name: str = 'FullyConnected', jit: bool = True, cuda_graphs: bool = True, amp: bool = True, amp_cpu: bool = None, amp_gpu: bool = None, torch_fx: bool = True, bf16: bool = False, onnx: bool = True, onnx_gpu: bool = None, onnx_cpu: bool = None, onnx_runtime: bool = True, trt: bool = False, var_dim: int = -1, func_torch: bool = True, auto_grad: bool = True)[source]

Fourier Neural Operators

class modulus.models.fno.fno.FNO(*args, **kwargs)[source]

Bases: Module

Fourier neural operator (FNO) model.

Note

The FNO architecture supports options for 1D, 2D, 3D and 4D fields which can be controlled using the dimension parameter.

Parameters

Example

Copy
Copied!

            
            >>> # define the 2d FNO model
>>> model = modulus.models.fno.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])

Note

Reference: Li, Zongyi, et al. “Fourier neural operator for parametric partial differential equations.” arXiv preprint arXiv:2010.08895 (2020).

forward(x: Tensor) → Tensor[source]

class modulus.models.fno.fno.FNO1DEncoder(in_channels: int = 1, num_fno_layers: int = 4, fno_layer_size: int = 32, num_fno_modes: Union[int, List[int]] = 16, padding: Union[int, List[int]] = 8, padding_type: str = 'constant', activation_fn: Module = GELU(approximate='none'), coord_features: bool = True)[source]

Bases: Module

1D Spectral encoder for FNO

Parameters

build_fno(num_fno_modes: List[int]) → None[source]

build_lift_network() → None[source]

forward(x: Tensor) → Tensor[source]

grid_to_points(value: Tensor) → Tuple[Tensor, List[int]][source]

converting from grid based (image) to point based representation

Parameters
Returns
Return type

meshgrid(shape: List[int], device: device) → Tensor[source]

Creates 1D meshgrid feature

Parameters
Returns
Return type

points_to_grid(value: Tensor, shape: List[int]) → Tensor[source]

converting from point based to grid based (image) representation

Parameters
Returns
Return type

class modulus.models.fno.fno.FNO2DEncoder(in_channels: int = 1, num_fno_layers: int = 4, fno_layer_size: int = 32, num_fno_modes: Union[int, List[int]] = 16, padding: Union[int, List[int]] = 8, padding_type: str = 'constant', activation_fn: Module = GELU(approximate='none'), coord_features: bool = True)[source]

Bases: Module

2D Spectral encoder for FNO

Parameters

build_fno(num_fno_modes: List[int]) → None[source]

build_lift_network() → None[source]

forward(x: Tensor) → Tensor[source]

grid_to_points(value: Tensor) → Tuple[Tensor, List[int]][source]

converting from grid based (image) to point based representation

Parameters
Returns
Return type

meshgrid(shape: List[int], device: device) → Tensor[source]

Creates 2D meshgrid feature

Parameters
Returns
Return type

points_to_grid(value: Tensor, shape: List[int]) → Tensor[source]

converting from point based to grid based (image) representation

Parameters
Returns
Return type

class modulus.models.fno.fno.FNO3DEncoder(in_channels: int = 1, num_fno_layers: int = 4, fno_layer_size: int = 32, num_fno_modes: Union[int, List[int]] = 16, padding: Union[int, List[int]] = 8, padding_type: str = 'constant', activation_fn: Module = GELU(approximate='none'), coord_features: bool = True)[source]

Bases: Module

3D Spectral encoder for FNO

Parameters

build_fno(num_fno_modes: List[int]) → None[source]

build_lift_network() → None[source]

forward(x: Tensor) → Tensor[source]

grid_to_points(value: Tensor) → Tuple[Tensor, List[int]][source]

converting from grid based (image) to point based representation

Parameters
Returns
Return type

meshgrid(shape: List[int], device: device) → Tensor[source]

Creates 3D meshgrid feature

Parameters
Returns
Return type

points_to_grid(value: Tensor, shape: List[int]) → Tensor[source]

converting from point based to grid based (image) representation

Parameters
Returns
Return type

class modulus.models.fno.fno.FNO4DEncoder(in_channels: int = 1, num_fno_layers: int = 4, fno_layer_size: int = 32, num_fno_modes: Union[int, List[int]] = 16, padding: Union[int, List[int]] = 8, padding_type: str = 'constant', activation_fn: Module = GELU(approximate='none'), coord_features: bool = True)[source]

Bases: Module

4D Spectral encoder for FNO

Parameters

build_fno(num_fno_modes: List[int]) → None[source]

build_lift_network() → None[source]

forward(x: Tensor) → Tensor[source]

grid_to_points(value: Tensor) → Tuple[Tensor, List[int]][source]

converting from grid based (image) to point based representation

Parameters
Returns
Return type

meshgrid(shape: List[int], device: device) → Tensor[source]

Creates 4D meshgrid feature

Parameters
Returns
Return type

points_to_grid(value: Tensor, shape: List[int]) → Tensor[source]

converting from point based to grid based (image) representation

Parameters
Returns
Return type

class modulus.models.fno.fno.MetaData(name: str = 'FourierNeuralOperator', jit: bool = True, cuda_graphs: bool = True, amp: bool = False, amp_cpu: bool = None, amp_gpu: bool = None, torch_fx: bool = False, bf16: bool = False, onnx: bool = False, onnx_gpu: bool = False, onnx_cpu: bool = False, onnx_runtime: bool = False, trt: bool = False, var_dim: int = 1, func_torch: bool = False, auto_grad: bool = False)[source]

class modulus.models.afno.afno.AFNO(*args, **kwargs)[source]

Bases: Module

Adaptive Fourier neural operator (AFNO) model.

Note

AFNO is a model that is designed for 2D images only.

Parameters

Example

Copy
Copied!

            
            >>> model = modulus.models.afno.AFNO(
...     inp_shape=[32, 32],
...     in_channels=2,
...     out_channels=1,
...     patch_size=(8, 8),
...     embed_dim=16,
...     depth=2,
...     num_blocks=2,
... )
>>> input = torch.randn(32, 2, 32, 32) #(N, C, H, W)
>>> output = model(input)
>>> output.size()
torch.Size([32, 1, 32, 32])

Note

Reference: Guibas, John, et al. “Adaptive fourier neural operators: Efficient token mixers for transformers.” arXiv preprint arXiv:2111.13587 (2021).

forward(x: Tensor) → Tensor[source]

forward_features(x: Tensor) → Tensor[source]

class modulus.models.afno.afno.AFNO2DLayer(hidden_size: int, num_blocks: int = 8, sparsity_threshold: float = 0.01, hard_thresholding_fraction: float = 1, hidden_size_factor: int = 1)[source]

Bases: Module

AFNO spectral convolution layer

Parameters

forward(x: Tensor) → Tensor[source]

class modulus.models.afno.afno.AFNOMlp(in_features: int, latent_features: int, out_features: int, activation_fn: Module = GELU(approximate='none'), drop: float = 0.0)[source]

Bases: Module

Fully-connected Multi-layer perception used inside AFNO

Parameters

forward(x: Tensor) → Tensor[source]

class modulus.models.afno.afno.Block(embed_dim: int, num_blocks: int = 8, mlp_ratio: float = 4.0, drop: float = 0.0, activation_fn: ~torch.nn.modules.module.Module = GELU(approximate='none'), norm_layer: ~torch.nn.modules.module.Module = <class 'torch.nn.modules.normalization.LayerNorm'>, double_skip: bool = True, sparsity_threshold: float = 0.01, hard_thresholding_fraction: float = 1.0)[source]

Bases: Module

AFNO block, spectral convolution and MLP

Parameters

forward(x: Tensor) → Tensor[source]

class modulus.models.afno.afno.MetaData(name: str = 'AFNO', jit: bool = False, cuda_graphs: bool = True, amp: bool = True, amp_cpu: bool = None, amp_gpu: bool = None, torch_fx: bool = False, bf16: bool = False, onnx: bool = False, onnx_gpu: bool = True, onnx_cpu: bool = False, onnx_runtime: bool = True, trt: bool = False, var_dim: int = 1, func_torch: bool = False, auto_grad: bool = False)[source]

class modulus.models.afno.afno.PatchEmbed(inp_shape: List[int], in_channels: int, patch_size: List[int] = [16, 16], embed_dim: int = 256)[source]

Bases: Module

Patch embedding layer

Converts 2D patch into a 1D vector for input to AFNO

Parameters

forward(x: Tensor) → Tensor[source]

Graph Neural Networks

class modulus.models.meshgraphnet.meshgraphnet.MeshGraphNet(*args, **kwargs)[source]

Bases: Module

MeshGraphNet network architecture

Parameters

Example

Copy
Copied!

            
            >>> model = modulus.models.meshgraphnet.MeshGraphNet(
...         input_dim_nodes=4,
...         input_dim_edges=3,
...         output_dim=2,
...     )
>>> graph = dgl.rand_graph(10, 5)
>>> node_features = torch.randn(10, 4)
>>> edge_features = torch.randn(5, 3)
>>> output = model(node_features, edge_features, graph)
>>> output.size()
torch.Size([10, 2])

Note

Reference: Pfaff, Tobias, et al. “Learning mesh-based simulation with graph networks.” arXiv preprint arXiv:2010.03409 (2020).

forward(node_features: Tensor, edge_features: Tensor, graph: Union[DGLGraph, List[DGLGraph], CuGraphCSC]) → Tensor[source]

class modulus.models.meshgraphnet.meshgraphnet.MeshGraphNetProcessor(processor_size: int = 15, input_dim_node: int = 128, input_dim_edge: int = 128, num_layers_node: int = 2, num_layers_edge: int = 2, aggregation: str = 'sum', norm_type: str = 'LayerNorm', activation_fn: Module = ReLU(), do_concat_trick: bool = False, num_processor_checkpoint_segments: int = 0)[source]

Bases: Module

MeshGraphNet processor block

forward(node_features: Tensor, edge_features: Tensor, graph: Union[DGLGraph, List[DGLGraph], CuGraphCSC]) → Tensor[source]

run_function(segment_start: int, segment_end: int) → Callable[[Tensor, Tensor, Union[DGLGraph, List[DGLGraph]]], Tuple[Tensor, Tensor]][source]

Custom forward for gradient checkpointing

Parameters
Returns
Return type

set_checkpoint_segments(checkpoint_segments: int)[source]

Set the number of checkpoint segments

Parameters
Raises

class modulus.models.meshgraphnet.meshgraphnet.MetaData(name: str = 'MeshGraphNet', jit: bool = False, cuda_graphs: bool = False, amp: bool = False, amp_cpu: bool = False, amp_gpu: bool = True, torch_fx: bool = False, bf16: bool = False, onnx: bool = False, onnx_gpu: bool = None, onnx_cpu: bool = None, onnx_runtime: bool = False, trt: bool = False, var_dim: int = -1, func_torch: bool = True, auto_grad: bool = True)[source]

class modulus.models.graphcast.graph_cast_net.GraphCastNet(*args, **kwargs)[source]

Bases: Module

GraphCast network architecture

Parameters

Note

Based on these papers:

“GraphCast: Learning skillful medium-range global weather forecasting”
https://arxiv.org/abs/2212.12794
“Forecasting Global Weather with Graph Neural Networks”
https://arxiv.org/abs/2202.07575
“Learning Mesh-Based Simulation with Graph Networks”
https://arxiv.org/abs/2010.03409
“MultiScale MeshGraphNets”
https://arxiv.org/abs/2210.00612

custom_forward(grid_nfeat: Tensor) → Tensor[source]

GraphCast forward method with support for gradient checkpointing.

Parameters
Returns
Return type

decoder_forward(mesh_efeat_processed: Tensor, mesh_nfeat_processed: Tensor, grid_nfeat_encoded: Tensor) → Tensor[source]

Forward method for the last layer of the processor, the decoder, and the final MLP.

Parameters
Returns
Return type

encoder_forward(grid_nfeat: Tensor) → Tensor[source]

Forward method for the embedder, encoder, and the first of the processor.

Parameters
Returns

forward(grid_nfeat: Tensor) → Tensor[source]

prepare_input(invar: Tensor, expect_partitioned_input: bool) → Tensor[source]

Prepares the input to the model in the required shape.

Parameters
Returns
Return type

prepare_output(outvar: Tensor, produce_aggregated_output: bool) → Tensor[source]

Prepares the output of the model in the shape [N, C, H, W].

Parameters
Returns
Return type

set_checkpoint_decoder(checkpoint_flag: bool)[source]

Sets checkpoint function for the last layer of the processor, the decoder, and the final MLP.

This function returns the appropriate checkpoint function based on the provided checkpoint_flag flag. If checkpoint_flag is True, the function returns the checkpoint function from PyTorch’s torch.utils.checkpoint. Otherwise, it returns an identity function that simply passes the inputs through the given layer.

Parameters
Returns
Return type

set_checkpoint_encoder(checkpoint_flag: bool)[source]

Sets checkpoint function for the embedder, encoder, and the first of the processor.

Parameters
Returns
Return type

set_checkpoint_model(checkpoint_flag: bool)[source]

Sets checkpoint function for the entire model.

This function returns the appropriate checkpoint function based on the provided checkpoint_flag flag. If checkpoint_flag is True, the function returns the checkpoint function from PyTorch’s torch.utils.checkpoint. In this case, all the other gradient checkpoitings will be disabled. Otherwise, it returns an identity function that simply passes the inputs through the given layer.

Parameters
Returns
Return type

set_checkpoint_processor(checkpoint_segments: int)[source]

Sets checkpoint function for the processor excluding the first and last layers.

This function returns the appropriate checkpoint function based on the provided checkpoint_segments flag. If checkpoint_segments is positive, the function returns the checkpoint function from PyTorch’s torch.utils.checkpoint, with number of checkpointing segments equal to checkpoint_segments. Otherwise, it returns an identity function that simply passes the inputs through the given layer.

Parameters
Returns
Return type

to(*args: Any, **kwargs: Any) → Self[source]

Moves the object to the specified device, dtype, or format. This method moves the object and its underlying graph and graph features to the specified device, dtype, or format, and returns the updated object.

Parameters
Returns
Return type

class modulus.models.graphcast.graph_cast_net.MetaData(name: str = 'GraphCastNet', jit: bool = False, cuda_graphs: bool = False, amp: bool = False, amp_cpu: bool = False, amp_gpu: bool = True, torch_fx: bool = False, bf16: bool = True, onnx: bool = False, onnx_gpu: bool = None, onnx_cpu: bool = None, onnx_runtime: bool = False, trt: bool = False, var_dim: int = -1, func_torch: bool = False, auto_grad: bool = False)[source]

Pix2Pix Net

class modulus.models.pix2pix.pix2pix.MetaData(name: str = 'Pix2Pix', jit: bool = True, cuda_graphs: bool = True, amp: bool = False, amp_cpu: bool = False, amp_gpu: bool = True, torch_fx: bool = False, bf16: bool = False, onnx: bool = True, onnx_gpu: bool = None, onnx_cpu: bool = None, onnx_runtime: bool = False, trt: bool = False, var_dim: int = 1, func_torch: bool = True, auto_grad: bool = True)[source]

class modulus.models.pix2pix.pix2pix.Pix2Pix(*args, **kwargs)[source]

Bases: Module

Convolutional encoder-decoder based on pix2pix generator models.

Note

The pix2pix architecture supports options for 1D, 2D and 3D fields which can be constroled using the dimension parameter.

Parameters

Example

Copy
Copied!

            
            >>> #2D convolutional encoder decoder
>>> model = modulus.models.pix2pix.Pix2Pix(
... in_channels=1,
... out_channels=2,
... dimension=2,
... conv_layer_size=4)
>>> input = torch.randn(4, 1, 32, 32) #(N, C, H, W)
>>> output = model(input)
>>> output.size()
torch.Size([4, 2, 32, 32])

Note

Reference: Isola, Phillip, et al. “Image-To-Image translation with conditional adversarial networks” Conference on Computer Vision and Pattern Recognition, 2017. https://arxiv.org/abs/1611.07004

Reference: Wang, Ting-Chun, et al. “High-Resolution image synthesis and semantic manipulation with conditional GANs” Conference on Computer Vision and Pattern Recognition, 2018. https://arxiv.org/abs/1711.11585

Note

Based on the implementation: https://github.com/NVIDIA/pix2pixHD

forward(input: Tensor) → Tensor[source]

class modulus.models.pix2pix.pix2pix.ResnetBlock(dimension: int, channels: int, padding_type: str = 'reflect', activation: Module = ReLU(), use_batch_norm: bool = False, use_dropout: bool = False)[source]

Bases: Module

A simple ResNet block

Parameters

forward(x: Tensor) → Tensor[source]

Recurrent Neural Networks

class modulus.models.rnn.rnn_one2many.MetaData(name: str = 'One2ManyRNN', jit: bool = False, cuda_graphs: bool = False, amp: bool = True, amp_cpu: bool = None, amp_gpu: bool = None, torch_fx: bool = True, bf16: bool = False, onnx: bool = False, onnx_gpu: bool = None, onnx_cpu: bool = None, onnx_runtime: bool = False, trt: bool = False, var_dim: int = -1, func_torch: bool = False, auto_grad: bool = False)[source]

class modulus.models.rnn.rnn_one2many.One2ManyRNN(*args, **kwargs)[source]

Bases: Module

A RNN model with encoder/decoder for 2d/3d problems that provides predictions based on single initial condition.

Parameters

Example

Copy
Copied!

            
            >>> model = modulus.models.rnn.One2ManyRNN(
... input_channels=6,
... dimension=2,
... nr_latent_channels=32,
... activation_fn="relu",
... nr_downsamples=2,
... nr_tsteps=16,
... )
>>> input = invar = torch.randn(4, 6, 1, 16, 16) # [N, C, T, H, W]
>>> output = model(input)
>>> output.size()
torch.Size([4, 6, 16, 16, 16])

forward(x: Tensor) → Tensor[source]

Forward pass

Parameters
Returns
Return type

class modulus.models.rnn.rnn_seq2seq.MetaData(name: str = 'Seq2SeqRNN', jit: bool = False, cuda_graphs: bool = False, amp: bool = True, amp_cpu: bool = None, amp_gpu: bool = None, torch_fx: bool = True, bf16: bool = False, onnx: bool = False, onnx_gpu: bool = None, onnx_cpu: bool = None, onnx_runtime: bool = False, trt: bool = False, var_dim: int = -1, func_torch: bool = False, auto_grad: bool = False)[source]

class modulus.models.rnn.rnn_seq2seq.Seq2SeqRNN(*args, **kwargs)[source]

Bases: Module

A RNN model with encoder/decoder for 2d/3d problems. Given input 0 to t-1, predicts signal t to t + nr_tsteps

Parameters

Example

Copy
Copied!

            
            >>> model = modulus.models.rnn.Seq2SeqRNN(
... input_channels=6,
... dimension=2,
... nr_latent_channels=32,
... activation_fn="relu",
... nr_downsamples=2,
... nr_tsteps=16,
... )
>>> input = invar = torch.randn(4, 6, 16, 16, 16) # [N, C, T, H, W]
>>> output = model(input)
>>> output.size()
torch.Size([4, 6, 16, 16, 16])

forward(x: Tensor) → Tensor[source]

Forward pass

Parameters
Returns
Return type

Super Resolution Network

class modulus.models.srrn.super_res_net.ConvolutionalBlock3d(in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, batch_norm: bool = False, activation_fn: Module = Identity())[source]

Bases: Module

3D convolutional block

Parameters

forward(input: Tensor) → Tensor[source]

class modulus.models.srrn.super_res_net.MetaData(name: str = 'SuperResolution', jit: bool = True, cuda_graphs: bool = False, amp: bool = False, amp_cpu: bool = False, amp_gpu: bool = False, torch_fx: bool = False, bf16: bool = False, onnx: bool = True, onnx_gpu: bool = None, onnx_cpu: bool = None, onnx_runtime: bool = False, trt: bool = False, var_dim: int = 1, func_torch: bool = True, auto_grad: bool = True)[source]

class modulus.models.srrn.super_res_net.PixelShuffle3d(scale: int)[source]

Bases: Module

3D pixel-shuffle operation

Parameters

Note

Reference: http://www.multisilicon.com/blog/a25332339.html

forward(input: Tensor) → Tensor[source]

class modulus.models.srrn.super_res_net.ResidualConvBlock3d(n_layers: int = 1, kernel_size: int = 3, conv_layer_size: int = 64, activation_fn: Module = Identity())[source]

Bases: Module

3D ResNet block

Parameters

forward(input: Tensor) → Tensor[source]

class modulus.models.srrn.super_res_net.SRResNet(*args, **kwargs)[source]

Bases: Module

3D convolutional super-resolution network

Parameters

Example

Copy
Copied!

            
            >>> #3D convolutional encoder decoder
>>> model = modulus.models.srrn.SRResNet(
... in_channels=1,
... out_channels=2,
... conv_layer_size=4,
... scaling_factor=2)
>>> input = torch.randn(4, 1, 8, 8, 8) #(N, C, D, H, W)
>>> output = model(input)
>>> output.size()
torch.Size([4, 2, 16, 16, 16])

Note

Based on the implementation: https://github.com/sgrvinod/a-PyTorch-Tutorial-to-Super-Resolution

forward(in_vars: Tensor) → Tensor[source]

class modulus.models.srrn.super_res_net.SubPixel_ConvolutionalBlock3d(kernel_size: int = 3, conv_layer_size: int = 64, scaling_factor: int = 2)[source]

Bases: Module

Convolutional block with Pixel Shuffle operation

Parameters

forward(input: Tensor) → Tensor[source]

DLWP Model

class modulus.models.dlwp.dlwp.DLWP(*args, **kwargs)[source]

Bases: Module

A Convolutional model for Deep Learning Weather Prediction that works on Cubed-sphere grids.

This model expects the input to be of shape [N, C, 6, Res, Res]

Parameters

Example

Copy
Copied!

            
            >>> model = modulus.models.dlwp.DLWP(
... nr_input_channels=2,
... nr_output_channels=4,
... )
>>> input = torch.randn(4, 2, 6, 64, 64) # [N, C, F, Res, Res]
>>> output = model(input)
>>> output.size()
torch.Size([4, 4, 6, 64, 64])

Note

Reference: Weyn, Jonathan A., et al. “Sub‐seasonal forecasting with a large ensemble

forward(cubed_sphere_input)[source]

class modulus.models.dlwp.dlwp.MetaData(name: str = 'DLWP', jit: bool = False, cuda_graphs: bool = True, amp: bool = False, amp_cpu: bool = True, amp_gpu: bool = True, torch_fx: bool = False, bf16: bool = False, onnx: bool = False, onnx_gpu: bool = None, onnx_cpu: bool = None, onnx_runtime: bool = False, trt: bool = False, var_dim: int = 1, func_torch: bool = False, auto_grad: bool = False)[source]

Previous Simple Logging and Checkpointing recipe

Next Modulus Datapipes