minerva.models.nets.unet

Full assembly of the parts to form the complete network

Classes

UNet

This class is a simple implementation of the U-Net model, which is a

_DoubleConv

(convolution => [BN] => ReLU) * 2

_Down

Downscaling with maxpool then double conv

_OutConv

_UNet

Implementation of U-Net model.

_Up

Upscaling then double conv

Module Contents

class minerva.models.nets.unet.UNet(n_channels=1, bilinear=False, learning_rate=0.001, loss_fn=None, **kwargs)

Bases: minerva.models.nets.base.SimpleSupervisedModel

This class is a simple implementation of the U-Net model, which is a convolutional neural network used for image segmentation. The model consists of a contracting path (encoder) and an expansive path (decoder). The contracting path follows the typical architecture of a convolutional neural network, with repeated applications of convolutions and max pooling layers. The expansive path consists of up-convolutions and concatenation of feature maps from the contracting path. The model also has skip connections, which allows the expansive path to use information from the contracting path at multiple resolutions. The U-Net model was originally proposed by Ronneberger, Fischer, and Brox in 2015.

This architecture, handles arbitrary input sizes, and returns an output of the same size as the input. The expected input size is (N, C, H, W), where N is the batch size, C is the number of channels, H is the height of the input image, and W is the width of the input image.

Note that, for this implementation, the input batch is a single tensor and not a tuple of tensors (e.g., data and label).

Note that this class wrappers the _UNet class, which is the actual implementation of the U-Net model, into a SimpleReconstructionNet class, which is a simple autoencoder pipeline for reconstruction tasks.

References

Ronneberger, Olaf, Philipp Fischer, and Thomas Brox. “U-net: Convolutional networks for biomedical image segmentation.” Medical Image Computing and Computer-Assisted Intervention-MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18. Springer International Publishing, 2015.

Wrapper implementation of the U-Net model.

Parameters

n_channelsint, optional

The number of channels of the input, by default 1

bilinearbool, optional

If True use bilinear interpolation for upsampling, by default False.

learning_ratefloat, optional

The learning rate to Adam optimizer, by default 1e-3

loss_fntorch.nn.Module, optional

The function used to compute the loss. If None, it will be used the MSELoss, by default None.

kwargsDict

Additional arguments to be passed to the SimpleSupervisedModel class.

Parameters:
  • n_channels (int)

  • bilinear (bool)

  • learning_rate (float)

  • loss_fn (Optional[torch.nn.Module])

class minerva.models.nets.unet._DoubleConv(in_channels, out_channels, mid_channels=None)

Bases: torch.nn.Module

(convolution => [BN] => ReLU) * 2

Parameters

in_channelsint

Number of input channels, i.e. the number of channels in the input image (1 for grayscale, 3 for RGB)

out_channelsint

Number of output channels, i.e. the number of channels produced by the convolution

mid_channelsint, optional

Number of channels in the middle, by default None

forward(x)
class minerva.models.nets.unet._Down(in_channels, out_channels)

Bases: torch.nn.Module

Downscaling with maxpool then double conv

forward(x)
class minerva.models.nets.unet._OutConv(in_channels, out_channels)

Bases: torch.nn.Module

forward(x)
class minerva.models.nets.unet._UNet(n_channels=1, bilinear=False)

Bases: torch.nn.Module

Implementation of U-Net model.

Implementation of U-Net model.

Parameters

n_channelsint, optional

Number of input channels, by default 1

bilinearbool, optional

If True use bilinear interpolation for upsampling, by default False.

forward(x)
Parameters:
  • n_channels (int)

  • bilinear (bool)

class minerva.models.nets.unet._Up(in_channels, out_channels, bilinear=True)

Bases: torch.nn.Module

Upscaling then double conv

forward(x1, x2)