mridc.collections.reconstruction.models.unet_base package

Submodules

mridc.collections.reconstruction.models.unet_base.unet_block module

class mridc.collections.reconstruction.models.unet_base.unet_block.ConvBlock(in_chans: int, out_chans: int, drop_prob: float)[source]

Bases: torch.nn.modules.module.Module

A Convolutional Block that consists of two convolution layers each followed by instance normalization, LeakyReLU activation and dropout.

forward(image: torch.Tensor) torch.Tensor[source]
Parameters

image (Input 4D tensor of shape (N, in_chans, H, W).) –

Return type

Output tensor of shape (N, out_chans, H, W).

training: bool
class mridc.collections.reconstruction.models.unet_base.unet_block.NormUnet(chans: int, num_pools: int, in_chans: int = 2, out_chans: int = 2, drop_prob: float = 0.0, padding_size: int = 15, normalize: bool = True, norm_groups: int = 2)[source]

Bases: torch.nn.modules.module.Module

Normalized U-Net model.

This is the same as a regular U-Net, but with normalization applied to the input before the U-Net. This keeps the values more numerically stable during training.

static chan_complex_to_last_dim(x: torch.Tensor) torch.Tensor[source]

Convert the last dimension of the input to complex.

static complex_to_chan_dim(x: torch.Tensor) torch.Tensor[source]

Convert the last dimension of the input to complex.

forward(x: torch.Tensor) torch.Tensor[source]

Forward pass of the network.

norm(x: torch.Tensor) Tuple[torch.Tensor, torch.Tensor, torch.Tensor][source]

Normalize the input.

pad(x: torch.Tensor) Tuple[torch.Tensor, Tuple[List[int], List[int], int, int]][source]

Pad the input with zeros to make it square.

training: bool
unnorm(x: torch.Tensor, mean: torch.Tensor, std: torch.Tensor) torch.Tensor[source]

Unnormalize the input.

static unpad(x: torch.Tensor, h_pad: List[int], w_pad: List[int], h_mult: int, w_mult: int) torch.Tensor[source]

Unpad the input.

class mridc.collections.reconstruction.models.unet_base.unet_block.TransposeConvBlock(in_chans: int, out_chans: int)[source]

Bases: torch.nn.modules.module.Module

A Transpose Convolutional Block that consists of one convolution transpose layers followed by instance normalization and LeakyReLU activation.

forward(image: torch.Tensor) torch.Tensor[source]
Parameters

image (Input 4D tensor of shape (N, in_chans, H, W).) –

Return type

Output tensor of shape (N, out_chans, H*2, W*2).

training: bool
class mridc.collections.reconstruction.models.unet_base.unet_block.Unet(in_chans: int, out_chans: int, chans: int = 32, num_pool_layers: int = 4, drop_prob: float = 0.0)[source]

Bases: torch.nn.modules.module.Module

PyTorch implementation of a U-Net model, as presented in 1.

References

1
  1. Ronneberger, P. Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In International Conference on Medical image computing and computer-assisted intervention, pages 234–241. Springer, 2015.

forward(image: torch.Tensor) torch.Tensor[source]
Parameters

image (Input 4D tensor of shape (N, in_chans, H, W).) –

Return type

Output tensor of shape (N, out_chans, H, W).

training: bool

Module contents