composer.models.unet.model#

The Unet architecture used in image segmentation. The example we are using is for BRATS medical brain tumor dataset.

See the Model Card for more details.

Classes

UNet

Unet Architecture adapted from NVidia Deep Learning Examples.

class composer.models.unet.model.UNet(in_channels, n_class, kernels, strides, normalization_layer, negative_slope, residual, dimension)[source]#

Bases: torch.nn.modules.module.Module

Unet Architecture adapted from NVidia Deep Learning Examples.

Parameters
  • in_channels (int) โ€“ Number of input channels.

  • n_class (int) โ€“ Number of output layers.

  • kernels (list) โ€“ Conv layer kernel sizes.

  • strides (list) โ€“ Conv layer strides.

  • normalization_layer (str) โ€“ Normalization layer type, one of ("batch", "instance").

  • negative_slope (float) โ€“ Leaky relu negative slope.

  • residual (bool) โ€“ Use residual connections.

  • dimension (int) โ€“ Filter dimensions.