mednet.models.segment.lwnet¶
Little W-Net (LWNET) network architecture, from [GAD+22].
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
|
Convolution block. |
|
ConvBridgeBlock implementation. |
|
Base little U-Net (LUNET) network architecture, from [GAD+22]. |
|
|
|
UpConvBlock implementation. |
|
Upsample block implementation. |
- class mednet.models.segment.lwnet.ConvBlock(input_channels, output_channels, kernel_size=3, add_conv2d=False, add_max_pool2d=True)[source]¶
Bases:
Module
Convolution block.
- Parameters:
input_channels – Number of input channels.
output_channels – Number of output channels.
kernel_size – Kernel Size.
add_conv2d – If True, adds a Conv2d layer.
add_max_pool2d – 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(input_channels, output_channels, upsampling_mode)[source]¶
Bases:
Module
Upsample block implementation.
- Parameters:
- 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, kernel_size)[source]¶
Bases:
Module
ConvBridgeBlock implementation.
- 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(input_channels, output_channels, kernel_size=3, upsampling_mode='up_conv', add_conv_bridge=False, add_conv2d=False)[source]¶
Bases:
Module
UpConvBlock implementation.
- Parameters:
input_channels (
int
) – Number of input channels.output_channels (
int
) – Number of output channels.kernel_size (
int
) – Kernel Size.upsampling_mode (
Literal
['transp_conv'
,'up_conv'
]) – Upsampling mode.add_conv_bridge (
bool
) – If True, adds a ConvBridgeBlock layer.add_conv2d (
bool
) – 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(input_channels, layers, num_classes=1, kernel_size=3, upsampling_mode='transp_conv', add_conv_bridge=True, add_conv2d=True)[source]¶
Bases:
Module
Base little U-Net (LUNET) network architecture, from [GAD+22].
- Parameters:
input_channels – Number of input channels.
layers – Number of layers of the model.
num_classes (
int
) – Number of outputs (classes) for this model.kernel_size (
int
) – Kernel Size.upsampling_mode (
Literal
['transp_conv'
,'up_conv'
]) – Upsampling mode.add_conv_bridge – If True, adds a ConvBridgeBlock layer.
add_conv2d – 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.losses.MultiLayerBCELogitsLossWeightedPerBatch'>, loss_arguments=None, optimizer_type=<class 'torch.optim.adam.Adam'>, optimizer_arguments=None, scheduler_type=None, scheduler_arguments=None, model_transforms=None, augmentation_transforms=None, num_classes=1)[source]¶
Bases:
Model
Little W-Net (LWNET) network architecture, from [GAD+22].
Concatenates two
Little U-Net
models.- Parameters:
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
] |None
) – Arguments to the loss.optimizer_type (
type
[Optimizer
]) – The type of optimizer to use for training.optimizer_arguments (
dict
[str
,Any
] |None
) – Arguments to the optimizer afterparams
.scheduler_type (
type
[LRScheduler
] |None
) – The type of scheduler to use for training.scheduler_arguments (
dict
[str
,Any
] |None
) – Arguments to the scheduler afterparams
.model_transforms (
Optional
[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 (
Optional
[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.
- property num_classes: int¶
Number of outputs (classes) for this model.
- Returns:
The number of outputs supported by this model.
- Return type:
- 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 callsforward()
. 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 forTrainer(strategy="ddp_spawn")
or training on 8 TPU cores withTrainer(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)