minerva.models.nets.unet
Full assembly of the parts to form the complete network
Classes
This class is a simple implementation of the U-Net model, which is a |
|
(convolution => [BN] => ReLU) * 2 |
|
Downscaling with maxpool then double conv |
|
Implementation of U-Net model. |
|
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)