mednet.models.segment.lwnet

Little W-Net (LWNET) network architecture, from [GALDRAN-2020].

It is based on two simple U-Nets with 3 layers concatenated to each other. The first U-Net produces a segmentation map that is used by the second to better guide segmentation.

Classes

ConvBlock(in_c, out_c[, k_sz, shortcut, pool])

Convolution block.

ConvBridgeBlock(channels[, k_sz])

ConvBridgeBlock implementation.

LittleUNet(in_c, n_classes, layers[, k_sz, ...])

Base little U-Net (LUNET) network architecture, from [GALDRAN-2020].

LittleWNet([loss_type, loss_arguments, ...])

Little W-Net (LWNET) network architecture, from [GALDRAN-2020].

UpConvBlock(in_c, out_c[, k_sz, up_mode, ...])

UpConvBlock implementation.

UpsampleBlock(in_c, out_c[, up_mode])

Upsample block implementation.
class mednet.models.segment.lwnet.ConvBlock(in_c, out_c, k_sz=3, shortcut=False, pool=True)[source]

Bases: Module

Convolution block.

Parameters:

  • in_c – Number of input channels.

  • out_c – Number of output channels.

  • k_sz – Kernel Size.

  • shortcut – If True, adds a Conv2d layer.

  • pool – If True, adds a MaxPool2d layer.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class mednet.models.segment.lwnet.UpsampleBlock(in_c, out_c, up_mode='transp_conv')[source]

Bases: Module

Upsample block implementation.

Parameters:

  • in_c – Number of input channels.

  • out_c – Number of output channels.

  • up_mode – Upsampling mode.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class mednet.models.segment.lwnet.ConvBridgeBlock(channels, k_sz=3)[source]

Bases: Module

ConvBridgeBlock implementation.

Parameters:

  • channels – Number of channels.

  • k_sz – Kernel Size.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class mednet.models.segment.lwnet.UpConvBlock(in_c, out_c, k_sz=3, up_mode='up_conv', conv_bridge=False, shortcut=False)[source]

Bases: Module

UpConvBlock implementation.

Parameters:

  • in_c – Number of input channels.

  • out_c – Number of output channels.

  • k_sz – Kernel Size.

  • up_mode – Upsampling mode.

  • conv_bridge – If True, adds a ConvBridgeBlock layer.

  • shortcut – If True, adds a Conv2d layer.

forward(x, skip)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class mednet.models.segment.lwnet.LittleUNet(in_c, n_classes, layers, k_sz=3, up_mode='transp_conv', conv_bridge=True, shortcut=True)[source]

Bases: Module

Base little U-Net (LUNET) network architecture, from [GALDRAN-2020].

Parameters:

  • in_c – Number of input channels.

  • n_classes – Number of outputs (classes) for this model.

  • layers – Number of layers of the model.

  • k_sz – Kernel Size.

  • up_mode – Upsampling mode.

  • conv_bridge – If True, adds a ConvBridgeBlock layer.

  • shortcut – If True, adds a Conv2d layer.

forward(x)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class mednet.models.segment.lwnet.LittleWNet(loss_type=<class 'mednet.models.segment.losses.MultiLayerBCELogitsLossWeightedPerBatch'>, loss_arguments={}, optimizer_type=<class 'torch.optim.adam.Adam'>, optimizer_arguments={}, scheduler_type=None, scheduler_arguments={}, model_transforms=[], augmentation_transforms=[], num_classes=1)[source]

Bases: Model

Little W-Net (LWNET) network architecture, from [GALDRAN-2020].

Concatenates two Little U-Net models.

Parameters:

  • loss_type (type[Module]) –

    The loss to be used for training and evaluation.

    Warning

    The loss should be set to always return batch averages (as opposed to the batch sum), as our logging system expects it so.

  • loss_arguments (dict[str, Any]) – Arguments to the loss.

  • optimizer_type (type[Optimizer]) – The type of optimizer to use for training.

  • optimizer_arguments (dict[str, Any]) – Arguments to the optimizer after params.

  • scheduler_type (type[LRScheduler] | None) – The type of scheduler to use for training.

  • scheduler_arguments (dict[str, Any]) – Arguments to the scheduler after params.

  • model_transforms (Sequence[Callable[[Tensor], Tensor]]) – An optional sequence of torch modules containing transforms to be applied on the input before it is fed into the network.

  • augmentation_transforms (Sequence[Callable[[Tensor], Tensor]]) – An optional sequence of torch modules containing transforms to be applied on the input before it is fed into the network.

  • num_classes (int) – Number of outputs (classes) for this model.

forward(x)[source]

Same as torch.nn.Module.forward().

Parameters:

  • *args – Whatever you decide to pass into the forward method.

  • **kwargs – Keyword arguments are also possible.

Returns:

Your model’s output

predict_step(batch, batch_idx, dataloader_idx=0)[source]

Step function called during predict(). By default, it calls forward(). Override to add any processing logic.

The predict_step() is used to scale inference on multi-devices.

To prevent an OOM error, it is possible to use BasePredictionWriter callback to write the predictions to disk or database after each batch or on epoch end.

The BasePredictionWriter should be used while using a spawn based accelerator. This happens for Trainer(strategy="ddp_spawn") or training on 8 TPU cores with Trainer(accelerator="tpu", devices=8) as predictions won’t be returned.

Parameters:

  • batch – The output of your data iterable, normally a DataLoader.

  • batch_idx – The index of this batch.

  • dataloader_idx – The index of the dataloader that produced this batch. (only if multiple dataloaders used)

Returns:

Predicted output (optional).

Example

class MyModel(LightningModule):

    def predict_step(self, batch, batch_idx, dataloader_idx=0):
        return self(batch)

dm = ...
model = MyModel()
trainer = Trainer(accelerator="gpu", devices=2)
predictions = trainer.predict(model, dm)