liom_toolkit.segmentation.vseg.training module

liom_toolkit.segmentation.vseg.training.create_images(x: Tensor, y: Tensor, pred: Tensor, num_images: int = 4) → list[ndarray]

Create images for visualization

Parameters:

  • x (torch.Tensor) – The input tensor

  • y (torch.Tensor) – The true labels

  • pred (torch.Tensor) – The predicted labels

  • num_images (int) – The number of images to create

Returns:

The images

Return type:

List[np.ndarray]

liom_toolkit.segmentation.vseg.training.evaluate(model: ~liom_toolkit.segmentation.vseg.model.VsegModel, loader: ~torch.utils.data.dataloader.DataLoader, loss_fn: ~torch.nn.modules.module.Module, device: ~torch.device) -> (<class 'float'>, <class 'torch.Tensor'>, <class 'torch.Tensor'>, <class 'torch.Tensor'>, <class 'float'>, <class 'float'>, <class 'float'>, <class 'float'>)

Evaluate the model for an epoch

Parameters:

  • model (VsegModel) – The model to evaluate

  • loader (torch.utils.data.DataLoader) – The data loader

  • loss_fn (torch.nn.Module) – The loss function

  • device (torch.device) – The device to use for evaluation

Returns:

The loss, the true labels, the predicted labels, the input, and the metrics

Return type:

(float, torch.Tensor, torch.Tensor, torch.Tensor, float, float, float, float)

liom_toolkit.segmentation.vseg.training.mask_image(x, y_mask, pred_mask, i)
liom_toolkit.segmentation.vseg.training.train(model: ~liom_toolkit.segmentation.vseg.model.VsegModel, loader: ~torch.utils.data.dataloader.DataLoader, optimizer: ~torch.optim.optimizer.Optimizer, loss_fn: ~torch.nn.modules.module.Module, device: ~torch.device) -> (<class 'float'>, <class 'torch.Tensor'>, <class 'torch.Tensor'>, <class 'torch.Tensor'>)

Train the model for an epoch

Parameters:

  • model (VsegModel) – The model to train

  • loader (torch.utils.data.DataLoader) – The data loader

  • optimizer (torch.optim.Optimizer) – The optimizer

  • loss_fn (torch.nn.Module) – The loss function

  • device (torch.device) – The device to use for training

Returns:

The loss, the true labels, the predicted labels, and the input

Return type:

(float, torch.Tensor, torch.Tensor, torch.Tensor)

liom_toolkit.segmentation.vseg.training.train_model(dataset_dir: str = 'data/patches', dev: str = 'cuda', output_train: str = 'data/training', learning_rate: float = 0.003673, batch_size: int = 35, epochs: int = 62) → None

Train the vessel segmentation model

Parameters:

  • dataset_dir (str) – The directory of the dataset

  • dev (str) – The device to use for training

  • output_train (str) – The output directory for the training

  • learning_rate (float) – The learning rate for the optimizer

  • batch_size (int) – The batch size for training

  • epochs (int) – The number of epochs to train

Returns:

None