Convolutional Layers

class physicsnemo.nn.module.conv_layers.Conv2d(

in_channels: int,

out_channels: int,

kernel: int,

bias: bool = True,

up: bool = False,

down: bool = False,

resample_filter: List[int] = [1, 1],

fused_resample: bool = False,

init_mode: Literal['xavier_uniform', 'xavier_normal', 'kaiming_uniform', 'kaiming_normal'] = 'kaiming_normal',

init_weight: float = 1.0,

init_bias: float = 0.0,

fused_conv_bias: bool = False,

amp_mode: bool = False,

)

Bases: Module

A custom 2D convolutional layer implementation with support for up-sampling, down-sampling, and custom weight and bias initializations. The layer’s weights and biases canbe initialized using custom initialization strategies like “kaiming_normal”, and can be further scaled by factors init_weight and init_bias.

Parameters:

• in_channels (int) – Number of channels in the input image.

• out_channels (int) – Number of channels produced by the convolution.

• kernel (int) – Size of the convolving kernel.

• bias (bool, optional) – The biases of the layer. If set to None, the layer will not learn an additive bias. By default True.

• up (bool, optional) – Whether to perform up-sampling. By default False.

• down (bool, optional) – Whether to perform down-sampling. By default False.

• resample_filter (List[int], optional) – Filter to be used for resampling. By default [1, 1].

• fused_resample (bool, optional) – If True, performs fused up-sampling and convolution or fused down-sampling and convolution. By default False.

• init_mode (str, optional (default="kaiming_normal")) – init_mode : str, optional (default=”kaiming_normal”) The mode/type of initialization to use for weights and biases. Supported modes are: - “xavier_uniform”: Xavier (Glorot) uniform initialization. - “xavier_normal”: Xavier (Glorot) normal initialization. - “kaiming_uniform”: Kaiming (He) uniform initialization. - “kaiming_normal”: Kaiming (He) normal initialization. By default “kaiming_normal”.

• init_weight (float, optional) – A scaling factor to multiply with the initialized weights. By default 1.0.

• init_bias (float, optional) – A scaling factor to multiply with the initialized biases. By default 0.0.

• fused_conv_bias (bool, optional) – A boolean flag indicating whether bias will be passed as a parameter of conv2d. By default False.

• amp_mode (bool, optional) – A boolean flag indicating whether mixed-precision (AMP) training is enabled. Defaults to False.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class physicsnemo.nn.module.conv_layers.ConvBlock(

in_chans,

out_chans,

num_groups,

num_residuals=2,

upsample=0,

)

Bases: Module

Conv2d block :param in_chans: Number of input channels. :type in_chans: int :param out_chans: Number of output channels. :type out_chans: int :param num_groups: Number of groups to separate the channels into for group normalization. :type num_groups: int :param num_residuals: Number of Conv2d operator. Default: 2 :type num_residuals: int, optinal :param upsample: 1: Upsample, 0: Conv, -1: Downsample. Default: 0 :type upsample: int, optinal

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class physicsnemo.nn.module.conv_layers.ConvGRULayer(*args, **kwargs)

Bases: Module

Convolutional GRU layer.

Parameters:

• in_features (int) – Input features/channels.

• hidden_size (int) – Hidden layer features/channels.

• dimension (int) – Spatial dimension of the input.

• activation_fn (nn.Module, optional, default=nn.ReLU()) – Activation Function to use.

Forward:

• x (torch.Tensor) – Input tensor of shape \((B, C_{in}, *)\) where \(*\) represents spatial dimensions.

• hidden (torch.Tensor) – Hidden state tensor of shape \((B, H, *)\) where \(H\) is hidden_size.

Outputs:

torch.Tensor – Next hidden state of shape \((B, H, *)\).

forward(

x: Float[Tensor, 'batch in_features ...'],

hidden: Float[Tensor, 'batch hidden_size ...'],

) → Float[Tensor, 'batch hidden_size ...']

Forward pass implementing GRU update.

class physicsnemo.nn.module.conv_layers.ConvLayer(*args, **kwargs)

Bases: Module

Generalized Convolution Block.

Parameters:

• in_channels (int) – Number of input channels.

• out_channels (int) – Number of output channels.

• dimension (int) – Dimensionality of the input (1, 2, or 3).

• kernel_size (int) – Kernel size for the convolution.

• stride (int, optional, default=1) – Stride for the convolution.

• activation_fn (nn.Module, optional, default=nn.Identity()) – Activation function to use.

Forward:

x (torch.Tensor) – Input tensor of shape \((B, C_{in}, *)\) where \(*\) represents spatial dimensions matching dimension.

Outputs:

torch.Tensor – Output tensor of shape \((B, C_{out}, *)\) where spatial dimensions depend on stride and padding.

forward(

x: Float[Tensor, 'batch in_channels ...'],

) → Float[Tensor, 'batch out_channels ...']

Forward pass with same padding.

class physicsnemo.nn.module.conv_layers.ConvResidualBlock(*args, **kwargs)

Bases: Module

Convolutional ResNet Block.

Parameters:

• in_channels (int) – Number of input channels.

• out_channels (int) – Number of output channels.

• dimension (int) – Dimensionality of the input.

• stride (int, optional, default=1) – Stride of the convolutions.

• gated (bool, optional, default=False) – Residual Gate activation.

• layer_normalization (bool, optional, default=False) – Whether to apply layer normalization.

• begin_activation_fn (bool, optional, default=True) – Whether to use activation function in the beginning.

• activation_fn (nn.Module, optional, default=nn.ReLU()) – Activation function to use.

Forward:

x (torch.Tensor) – Input tensor of shape \((B, C_{in}, *)\) where \(*\) represents spatial dimensions matching dimension.

Outputs:

torch.Tensor – Output tensor of shape \((B, C_{out}, *)\) with residual connection.

Raises:

ValueError – If stride > 2 (not supported).

forward(

x: Float[Tensor, 'batch in_channels ...'],

) → Float[Tensor, 'batch out_channels ...']

Forward pass with residual connection.

class physicsnemo.nn.module.conv_layers.CubeEmbedding(

img_size,

patch_size,

in_chans,

embed_dim,

norm_layer=<class 'torch.nn.modules.normalization.LayerNorm'>,

)

Bases: Module

3D Image Cube Embedding :param img_size: Image size [T, Lat, Lon]. :type img_size: tuple[int] :param patch_size: Patch token size [T, Lat, Lon]. :type patch_size: tuple[int] :param in_chans: Number of input image channels. :type in_chans: int :param embed_dim: Number of projection output channels. :type embed_dim: int :param norm_layer: Normalization layer. Default: torch.nn.LayerNorm :type norm_layer: nn.Module, optional

forward(x: Tensor)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class physicsnemo.nn.module.conv_layers.TransposeConvLayer(*args, **kwargs)

Bases: Module

Generalized Transposed Convolution Block.

Parameters:

• in_channels (int) – Number of input channels.

• out_channels (int) – Number of output channels.

• dimension (int) – Dimensionality of the input (1, 2, or 3).

• kernel_size (int) – Kernel size for the convolution.

• stride (int, optional, default=1) – Stride for the convolution.

• activation_fn (nn.Module, optional, default=nn.Identity()) – Activation function to use.

Forward:

x (torch.Tensor) – Input tensor of shape \((B, C_{in}, *)\) where \(*\) represents spatial dimensions matching dimension.

Outputs:

torch.Tensor – Output tensor of shape \((B, C_{out}, *)\) where spatial dimensions are upsampled based on stride.

forward(

x: Float[Tensor, 'batch in_channels ...'],

) → Float[Tensor, 'batch out_channels ...']

Forward pass with transposed convolution and cropping.