Modulus Utils

Bases: _StaticCapture

An performance optimization decorator for PyTorch no grad evaluation.

This class should be initialized as a decorator on a function that computes run the forward pass of the model that does not require gradient calculations. This is the recommended method to use for inference and validation methods.

Parameters

• model (modulus.models.Module) – Modulus Model

• logger (Union[Logger, None], optional) – Modulus Launch Logger, by default None

• use_graphs (bool, optional) – Toggle CUDA graphs if supported by model, by default True

• use_amp (bool, optional) – Toggle AMP if supported by mode, by default True

• cuda_graph_warmup (int, optional) – Number of warmup steps for cuda graphs, by default 11

• amp_type (Union[float16, bfloat16], optional) – Auto casting type for AMP, by default torch.float16

• label (Optional[str, None], optional) – Static capture checkpoint label, by default None

Raises

ValueError – If the model provided is not a modulus.models.Module. I.e. has no meta data.

Example

            

            
>>> # Create model
>>> model = modulus.models.mlp.FullyConnected(2, 64, 2)
>>> input = torch.rand(8, 2)
>>> # Create evaluate function with optimization wrapper
>>> @StaticCaptureEvaluateNoGrad(model=model)
... def eval_step(model, invar):
...     predvar = model(invar)
...     return predvar
...
>>> output = eval_step(model, input)
>>> output.size()
torch.Size([8, 2])
        

Bases: _StaticCapture

A performance optimization decorator for PyTorch training functions.

This class should be initialized as a decorator on a function that computes the forward pass of the neural network and loss function. The user should only call the defind training step function. This will apply optimizations including: AMP and Cuda Graphs.

Parameters

• model (modulus.models.Module) – Modulus Model

• optim (torch.optim) – Optimizer

• logger (Union[Logger, None], optional) – Modulus Launch Logger, by default None

• use_graphs (bool, optional) – Toggle CUDA graphs if supported by model, by default True

• use_amp (bool, optional) – Toggle AMP if supported by mode, by default True

• cuda_graph_warmup (int, optional) – Number of warmup steps for cuda graphs, by default 11

• amp_type (Union[float16, bfloat16], optional) – Auto casting type for AMP, by default torch.float16

• label (Optional[str, None], optional) – Static capture checkpoint label, by default None

Raises

ValueError – If the model provided is not a modulus.models.Module. I.e. has no meta data.

Example

            

            
>>> # Create model
>>> model = modulus.models.mlp.FullyConnected(2, 64, 2)
>>> input = torch.rand(8, 2)
>>> output = torch.rand(8, 2)
>>> # Create optimizer
>>> optim = torch.optim.Adam(model.parameters(), lr=0.001)
>>> # Create training step function with optimization wrapper
>>> @StaticCaptureTraining(model=model, optim=optim)
... def training_step(model, invar, outvar):
...     predvar = model(invar)
...     loss = torch.sum(torch.pow(predvar - outvar, 2))
...     return loss
...
>>> # Sample training loop
>>> for i in range(3):
...     loss = training_step(model, input, output)
...
        

Bases: object

Class to load static data from netCDF files. Static data includes land-sea mask, geopotential, and latitude-longitude coordinates.

Parameters

• static_dataset_path (str) – Path to directory containing static data.

• latitudes (Tensor) – Tensor with shape (lat,) that includes latitudes.

• longitudes (Tensor) – Tensor with shape (lon,) that includes longitudes.

get() → Tensor[source]

Get all static data.

Returns

Tensor with shape (1, 5, lat, lon) that includes land-sea mask, geopotential, cosine of latitudes, sine and cosine of longitudes.

Return type

Tensor

get_geop(normalize: bool = True) → Tensor[source]

Get geopotential from netCDF file.

Parameters

normalize (bool, optional) – Whether to normalize the geopotential, by default True

Returns

Normalized geopotential with shape (1, 1, lat, lon).

Return type

Tensor

get_lat_lon() → Tensor[source]

Computes cosine of latitudes and sine and cosine of longitudes.

Returns

Tensor with shape (1, 3, lat, lon) tha includes cosine of latitudes, sine and cosine of longitudes.

Return type

Tensor

get_lsm() → Tensor[source]

Get land-sea mask from netCDF file.

Returns

Land-sea mask with shape (1, 1, lat, lon).

Return type

Tensor

Bases: object

Graph class for creating the graph2mesh, multimesh, and mesh2graph graphs.

Parameters

• icospheres_path (str) – Path to the icospheres json file. If the file does not exist, it will try to generate it using PyMesh.

• lat_lon_grid (Tensor) – Tensor with shape (lat, lon, 2) that includes the latitudes and longitudes meshgrid.

• dtype (torch.dtype, optional) – Data type of the graph, by default torch.float

create_g2m_graph(verbose: bool = True) → Tensor[source]

Create the graph2mesh graph.

Parameters

verbose (bool, optional) – verbosity, by default True

Returns

Graph2mesh graph.

Return type

DGLGraph

create_m2g_graph(verbose: bool = True) → Tensor[source]

Create the mesh2grid graph.

Parameters

verbose (bool, optional) – verbosity, by default True

Returns

Mesh2grid graph.

Return type

DGLGraph

create_mesh_graph(verbose: bool = True) → Tensor[source]

Create the multimesh graph.

Parameters

verbose (bool, optional) – verbosity, by default True

Returns

Multimesh graph.

Return type

DGLGraph

Bases: Module

Loss function with cell area weighting.

Parameters

area (torch.Tensor) – Cell area with shape [H, W].

forward(invar, outvar)[source]

Implicit forward function which computes the loss given a prediction and the corresponding targets.

Parameters

• invar (torch.Tensor) – prediction of shape [T, C, H, W].

• outvar (torch.Tensor) – target values of shape [T, C, H, W].

Bases: Function

Autograd fuunction for custom loss with cell area weighting.

static backward(ctx, grad_loss: Tensor)[source]

Backward method of custom loss function with cell area weighting.

static forward(ctx, invar: Tensor, outvar: Tensor, area: Tensor)[source]

Forward of custom loss function with cell area weighting.

Bases: CellAreaWeightedLossFunction

Custom loss function with cell area weighting.

Parameters

area (torch.Tensor) – Cell area with shape [H, W].

forward(invar: Tensor, outvar: Tensor) → Tensor[source]

Implicit forward function which computes the loss given a prediction and the corresponding targets.

Parameters

• invar (torch.Tensor) – prediction of shape [T, C, H, W].

• outvar (torch.Tensor) – target values of shape [T, C, H, W].

Bases: object

A package

Represents a potentially remote directory tree

get(path: str, recursive: bool = False) → str[source]

Get a local path to the item at path

path might be a remote file, in which case it is downloaded to a local cache at $LOCAL_CACHE or $HOME/.cache/modulus first.