Convolutional Networks#

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

in_channels (int) – Number of input channels
out_channels (Union[int, Any], optional) – Number of output channels
dimension (int) – Model dimensionality (supports 1, 2, 3).
conv_layer_size (int, optional) – Latent channel size after first convolution, by default 64
n_downsampling (int, optional) – Number of downsampling blocks, by default 3
n_upsampling (int, optional) – Number of upsampling blocks, by default 3
n_blocks (int, optional) – Number of residual blocks in middle of model, by default 3
activation_fn (Any, optional) – Activation function, by default “relu”
batch_norm (bool, optional) – Batch normalization, by default False
padding_type (str, optional) – Padding type (‘reflect’, ‘replicate’ or ‘zero’), by default “reflect”

Example

>>> #2D convolutional encoder decoder
>>> model = physicsnemo.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: NVIDIA/pix2pixHD

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

dimension (int) – Model dimensionality (supports 1, 2, 3).
channels (int) – Number of feature channels
padding_type (str, optional) – Padding type (‘reflect’, ‘replicate’ or ‘zero’), by default “reflect”
activation (nn.Module, optional) – Activation function, by default nn.ReLU()
use_batch_norm (bool, optional) – Batch normalization, by default False

class physicsnemo.models.pix2pix.pix2pixunet.Pix2PixUnet(*args, **kwargs)[source]#

Bases: Module

Convolutional encoder-decoder based on pix2pix generator models using Unet.

Note

The pix2pix with Unet architecture only supports 2D field.

Parameters:

in_channels (int) – Number of input channels
out_channels (Union[int, Any], optional) – Number of output channels
n_downsampling (int) – Number of downsampling in UNet
filter_size (int, optional) – Number of filters in last convolution layer, by default 64
norm_layer (optional) – Normalization layer, by default nn.BatchNorm2d
use_dropout (bool, optional) – Use dropout layers, by default False

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: junyanz/pytorch-CycleGAN-and-pix2pix

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

Bases: Module

3D convolutional super-resolution network

Parameters:

in_channels (int) – Number of input channels
out_channels (int) – Number of outout channels
large_kernel_size (int, optional) – convolutional kernel size for first and last convolution, by default 7
small_kernel_size (int, optional) – convolutional kernel size for internal convolutions, by default 3
conv_layer_size (int, optional) – Latent channel size, by default 32
n_resid_blocks (int, optional) – Number of residual blocks before , by default 8
scaling_factor (int, optional) – Scaling factor to increase the output feature size compared to the input (2, 4, or 8), by default 8
activation_fn (Any, optional) – Activation function, by default “prelu”

Example

>>> #3D convolutional encoder decoder
>>> model = physicsnemo.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: sgrvinod/a-PyTorch-Tutorial-to-Super-Resolution

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

in_channels (int) – Input channels
out_channels (int) – Output channels
kernel_size (int) – Kernel size
stride (int, optional) – Convolutional stride, by default 1
batch_norm (bool, optional) – Use batchnorm, by default False

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

Bases: Module

3D pixel-shuffle operation

Parameters:: scale (int) – Factor to downscale channel count by

Note

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

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

n_layers (int, optional) – Number of convolutional layers, by default 1
kernel_size (int, optional) – Kernel size, by default 3
conv_layer_size (int, optional) – Latent channel size, by default 64
activation_fn (nn.Module, optional) – Activation function, by default nn.Identity()

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

kernel_size (int, optional) – Kernel size, by default 3
conv_layer_size (int, optional) – Latent channel size, by default 64
scaling_factor (int, optional) – Pixel shuffle scaling factor, by default 2