mednet.models.classify.model

Definition of base model type for classification tasks.

Classes

Model(name[, training_step_fn, loss_type, ...])

Base model type for classification tasks.
class mednet.models.classify.model.Model(name, training_step_fn=default_training_step, loss_type=None, loss_arguments=None, optimizer_type=torch.optim.Adam, optimizer_arguments=None, scheduler_type=None, scheduler_arguments=None, model_transforms=None, augmentation_transforms=None, num_classes=1)[source]

Bases: Model

Base model type for classification tasks.

Parameters:

  • name (str) – Common name to give to models of this type.

  • training_step_fn (Callable) – The training step function to use.

  • loss_type (type[Module] | None) –

    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 after params.

  • scheduler_type (type[LRScheduler] | None) – The type of scheduler to use for training.

  • scheduler_arguments (dict[str, Any] | None) – Arguments to the scheduler after params.

  • 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.

training_step(batch, batch_idx)[source]

Perform a training step.

Parameters:

  • batch – The batch to apply the training step on. Should contain image and target keys.

  • batch_idx – The index of the batch, will be ignored.

Returns:

Training loss for the batch.

Return type:

torch.Tensor