Train Network

class train.MainModule

Bases: object

Class that defines the main module that will be used to train, test and predict with different medical models.

predict(args, log_dir, device, root_dir='.')

Function to predict with the model. Call the specific dataset and model, and predict with the model.

Parameters:

• args (argparse.Namespace) – Arguments from the command line.

• log_dir (str) – Log directory

• device (torch.device) – Device to be used

• root_dir (str, optional) – Root path dir. Defaults to “.”.

test(args, log_dir, device, root_dir='.')

Function to test the model. Call the specific dataset and model, and test the model.

Parameters:

• args (argparse.Namespace) – Arguments from the command line.

• log_dir (str) – Log directory

• device (torch.device) – Device to be used

• root_dir (str, optional) – Root path dir. Defaults to “.”.

train(args, log_dir)

Function to train the model. Call the specific dataset and model, and train the model.

Parameters:

• args (argparse.Namespace) – Arguments from the command line.

• log_dir (str) – Log directory

class train.Net(args)

Bases: LightningModule

Class that defines the Lightning Module that will be used for training, validation and testing.

configure_optimizers()

Function that configures the optimizer to be used during training.

Returns:

Optimizer to be used during training

Return type:

torch.optim.adamw.AdamW

forward(x)

Function that performs a forward pass on the network.

Parameters:

x (torch.Tensor | monai.data.meta_tensor.MetaTensor) – Input data to the network

Returns:

Output data from the network

Return type:

(torch.Tensor | monai.data.meta_tensor.MetaTensor)

on_test_epoch_end()

Function that performs an action at the end of the test epoch.

Returns:

Dictionary containing the tensorboard logs

Return type:

dict

on_training_epoch_end()

Function that performs an action at the end of the training epoch.

on_validation_epoch_end()

Function that performs an action at the end of the validation epoch.

Returns:

Dictionary containing the tensorboard logs

Return type:

dict

test_step(batch, batch_idx)

Function that performs the test step on the network.

Parameters:

• batch (dict) – The batch of data to be used for training

• batch_idx (int) – The index of the batch

Returns:

Dictionary containing the loss and the tensorboard logs

Return type:

dict

training_step(batch, batch_idx)

Function that performs a training step on the network.

Parameters:

• batch (dict) – The batch of data to be used for training

• batch_idx (int) – The index of the batch

Returns:

Dictionary containing the loss and the tensorboard logs

Return type:

dict

validation_step(batch, batch_idx)

Function that performs a validation step on the network.

Parameters:

• batch (dict) – The batch of data to be used for training

• batch_idx (int) – The index of the batch

Returns:

Dictionary containing the loss and the tensorboard logs

Return type:

dict

train.get_parser()

Function to get the parser with the arguments.

Raises:

ValueError – The specified configuration doesn’t exist

Returns:

Arguments from the command line.

Return type:

argparse.Namespace