biapy.engine.self_supervised

Self-supervised workflow for BiaPy.

This module defines the Self_supervised_Workflow class, which implements the training, validation, and inference pipeline for self-supervised learning tasks in BiaPy. It supports pretext tasks such as masking and crappification, and handles data preparation, model setup, metrics, predictions, and result saving for learning representations without labels.

class biapy.engine.self_supervised.Self_supervised_Workflow(cfg, job_identifier, device, system_dict, args, **kwargs)[source]

Bases: Base_Workflow

Self supervised workflow where the goal is to pretrain the backbone model by solving a so-called pretext task without labels. This way, the model learns a representation that can be later transferred to solve a downstream task in a labeled (but smaller) dataset.

More details in our documentation.

Parameters:

cfg (YACS configuration) – Running configuration.
Job_identifier (str) – Complete name of the running job.
device (Torch device) – Device used.
args (argpase class) – Arguments used in BiaPy’s call.

define_activations_and_channels()[source]

Define the activations to be applied to the model output and the channels that the model will output.

This function must define the following variables:

self.model_output_channelsList of int: Number of channels for each output head of the model. E.g. [3] for a model with one head outputting 3 channels, [1, 5] for a model with two heads outputting 1 and 5 channels respectively, etc.
self.model_output_channel_infoList of str: Information about the output channels. A value per output head of the model must be defined.
self.separated_class_channelbool: Whether if we should expect a separated output channel for classification.
self.head_activationsList of str: Activations to be applied to the model output. A value per output channel (not output head) of the model must be defined. “linear” and “ce_sigmoid” will not be applied. E.g. [“linear”] for a model with one channel, [“linear”, “sigmoid”] for a model with two channels, etc.

Example of a correct definition of the function for a model with two output heads: 1) the first one will be predicting foreground and contours; 2) the second one will classify into 3 classes the predicted objects. In this case the following definition would be correct:

self.model_output_channels = [1, 3]
self.model_output_channel_info = ["mask", "class"]
self.separated_class_channel = True
self.head_activations = ["ce_sigmoid", "ce_sigmoid", "ce_softmax", "ce_softmax", "ce_softmax"]

define_metrics()[source]

Define the metrics to be used in the workflow.

This function must define the following variables:

self.train_metricsList of functions: Metrics to be calculated during model’s training.
self.train_metric_namesList of str: Names of the metrics calculated during training.
self.train_metric_bestList of str: To know which value should be considered as the best one. Options must be: “max” or “min”.
self.test_metricsList of functions: Metrics to be calculated during model’s test/inference.
self.test_metric_namesList of str: Names of the metrics calculated during test/inference.
self.lossFunction: Loss function used during training and test.

MaskedAutoencoderViT_loss_wrapper(output, targets)[source]: Unravel MAE loss.

metric_calculation(output: ndarray[tuple[int, ...], dtype[_ScalarType_co]] | Tensor, targets: ndarray[tuple[int, ...], dtype[_ScalarType_co]] | Tensor, train: bool = True, metric_logger: MetricLogger | None = None) → Dict[source]

Calculate the metrics defined in define_metrics() function.

Parameters:

output (Torch Tensor) – Prediction of the model.
targets (Torch Tensor) – Ground truth to compare the prediction with.
train (bool, optional) – Whether to calculate train or test metrics.
metric_logger (MetricLogger, optional) – Class to be updated with the new metric(s) value(s) calculated.

Returns:

out_metrics – Value of the metrics for the given prediction.

Return type:

dict

prepare_targets(targets, batch)[source]

Perform any necessary data transformations to targets before calculating the loss.

Parameters:

targets (Torch Tensor) – Ground truth to compare the prediction with.
batch (Torch Tensor) – Prediction of the model.

Returns:

targets – Resulting targets.

Return type:

Torch tensor

process_test_sample()[source]: Process a sample in the test/inference phase.

torchvision_model_call(in_img: Tensor, is_train: bool = False) → Tensor | None[source]

Call a regular Pytorch model.

Parameters:

in_img (torch.Tensor) – Input image to pass through the model.
is_train (bool, optional) – Whether if the call is during training or inference.

Returns:

prediction – Image prediction.

Return type:

torch.Tensor

after_merge_patches(pred)[source]

Execute steps needed after merging all predicted patches into the original image.

Parameters:: pred (Torch Tensor) – Model prediction.

after_full_image(pred: ndarray[tuple[int, ...], dtype[_ScalarType_co]])[source]

Execute steps needed after generating the prediction by supplying the entire image to the model.

Parameters:: pred (NDArray) – Model prediction.

after_all_images()[source]: Execute steps needed after predicting all images.

prepare_ssl_data()[source]

Create self supervised “ground truth” images, if crappify was selected, to train the model based on the input images provided.

Images will be saved in a separate folder in the root path of the input images.

after_all_chunk_prediction_workflow_process()[source]: Place any code that needs to be done after predicting all patches in “by chunks” setting. This function is called on all ranks.

after_all_chunk_prediction_workflow_process_master_rank()[source]: Place any code that needs to be done after predicting all patches in “by chunks” setting, but only on the master rank. This function is called only on the master rank.