Recurrent Neural Networks#

class physicsnemo.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:

input_channels (int) – Number of channels in the input.
dimension (int, optional, default=2) – Spatial dimension of the input. Only 2D and 3D are supported.
nr_latent_channels (int, optional, default=512) – Channels for encoding/decoding.
nr_residual_blocks (int, optional, default=2) – Number of residual blocks.
activation_fn (str, optional, default="relu") – Activation function to use.
nr_downsamples (int, optional, default=2) – Number of downsamples.
nr_tsteps (int, optional, default=32) – Time steps to predict.

Forward:

x (torch.Tensor) – Input tensor of shape \((N, C, 1, H, W)\) for 2D or \((N, C, 1, D, H, W)\) for 3D, where \(N\) is the batch size, \(C\) is the number of channels, 1 is the number of input timesteps, and \(D, H, W\) are spatial dimensions.

Outputs:

torch.Tensor – Output tensor of shape \((N, C, T, H, W)\) for 2D or \((N, C, T, D, H, W)\) for 3D, where \(T\) is the number of timesteps being predicted.

Examples

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

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