minerva.models.nets.image.wisenet

Classes

WiseNet	A modular Lightning model wrapper for supervised learning tasks.
_WiseNet	Base class for all neural network modules.

Module Contents

class minerva.models.nets.image.wisenet.WiseNet(in_channels=1, out_channels=1, loss_fn=None, learning_rate=0.001, **kwargs)

Bases: minerva.models.nets.base.SimpleSupervisedModel

A modular Lightning model wrapper for supervised learning tasks.

This class enables the construction of supervised models by combining a backbone (feature extractor), an optional adapter, and a fully connected (FC) head. It provides a clean interface for setting up custom training, validation, and testing pipelines with pluggable loss functions, metrics, optimizers, and learning rate schedulers.

The architecture is structured as follows:

Backbone Model

v

Adapter (Optional)

(Flatten if needed)

v

Fully Connected Head

v

Loss Function

Training and validation steps comprise the following steps:

1. Forward pass input through the backbone.

2. Pass through adapter (if provided).

3. Flatten the output (if flatten is True) before the FC head.

4. Forward through the FC head.

5. Compute loss with respect to targets.

6. Backpropagate and update parameters.

7. Compute metrics and log them.

8. Return loss. train_loss, val_loss, and test_loss are always logged, along with any additional metrics specified in the train_metrics, val_metrics, and test_metrics dictionaries.

This wrapper is especially useful to quickly set up supervised models for various tasks, such as image classification, object detection, and segmentation. It is designed to be flexible and extensible, allowing users to easily swap out components like the backbone, adapter, and FC head as needed. The model is built with a focus on simplicity and modularity, making it easy to adapt to different use cases and requirements. The model is designed to be used with PyTorch Lightning and is compatible with its training loop.

Note: For more complex architectures that does not follow the above structure should not inherit from this class.

Note: Input batches must be tuples (input_tensor, target_tensor).

Initializes the supervised model with training components and configs.

Parameters

backbonetorch.nn.Module or LoadableModule

The backbone (feature extractor) model.

fctorch.nn.Module or LoadableModule

The fully connected head. Use nn.Identity() if not required.

loss_fntorch.nn.Module

Loss function to optimize during training.

adapterCallable, optional

Function to transform backbone outputs before feeding into fc.

learning_ratefloat, default=1e-3

Learning rate used for optimization.

flattenbool, default=True

If True, flattens backbone outputs before fc.

train_metricsdict, optional

TorchMetrics dictionary for training evaluation.

val_metricsdict, optional

TorchMetrics dictionary for validation evaluation.

test_metricsdict, optional

TorchMetrics dictionary for test evaluation.

freeze_backbonebool, default=False

If True, backbone parameters are frozen during training.

optimizer: type

Optimizer class to be instantiated. By default, it is set to torch.optim.Adam. Should be a subclass of torch.optim.Optimizer (e.g., torch.optim.SGD).

optimizer_kwargsdict, optional

Additional kwargs passed to the optimizer constructor.

lr_schedulertype, optional

Learning rate scheduler class to be instantiated. By default, it is set to None, which means no scheduler will be used. Should be a subclass of torch.optim.lr_scheduler.LRScheduler (e.g., torch.optim.lr_scheduler.StepLR).

lr_scheduler_kwargsdict, optional

Additional kwargs passed to the scheduler constructor.

_single_step(batch, batch_idx, step_name)

Perform a single train/validation/test step. It consists in making a forward pass with the input data on the backbone model, computing the loss between the output and the input data, and logging the loss.

Parameters

batchtorch.Tensor

The input data. It must be a 2-element tuple of tensors, where the first tensor is the input data and the second tensor is the mask.

batch_idxint

The index of the batch.

step_namestr

The name of the step. It will be used to log the loss. The possible values are: “train”, “val” and “test”. The loss will be logged as “{step_name}_loss”.

Returns

torch.Tensor

A tensor with the loss value.

Parameters:

• batch (torch.Tensor)

• batch_idx (int)

• step_name (str)

Return type:

torch.Tensor

predict_step(batch, batch_idx, dataloader_idx=None)

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.

Args:

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)

Return:

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)

Parameters:

• in_channels (int)

• out_channels (int)

• loss_fn (torch.nn.Module)

• learning_rate (float)

class minerva.models.nets.image.wisenet._WiseNet(in_channels=1, out_channels=1)

Bases: torch.nn.Module

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:

training (bool) – Boolean represents whether this module is in training or evaluation mode.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

conv1

conv2

conv3

conv4

conv5

conv6

conv7

conv8

forward(x)

in_channels = 1

out_channels = 1

pool1

pool2

pool3

pool4

relu