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 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.