mednet.models.temporal.model

Definition of base model type for temporal classification tasks.

Classes

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

Base model type for temporal classification tasks.

class mednet.models.temporal.model.Model(name, 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, task_type=None)

Bases: Model

Base model type for temporal classification tasks.

This class keeps the same configuration interface as classification models, but assumes each batch carries a lengths key with sequence lengths.

Validation logs loss/validation and auc/validation (macro AUROC for multiclass, same idea as mednet.models.classify.vit.ViT extensions).

Parameters:

• name (str) – Human-readable model name used for logging.

• loss_type (type[Module] | None) – Loss module type; defaults to torch.nn.BCEWithLogitsLoss.

• loss_arguments (dict[str, Any] | None) – Keyword arguments forwarded to loss_type.

• optimizer_type (type[Optimizer]) – Optimizer class; defaults to torch.optim.Adam.

• optimizer_arguments (dict[str, Any] | None) – Keyword arguments forwarded to optimizer_type (after params).

• scheduler_type (type[LRScheduler] | None) – Optional LR scheduler class.

• scheduler_arguments (dict[str, Any] | None) – Keyword arguments forwarded to scheduler_type (after optimizer).

• model_transforms (Optional[Sequence[Callable[[Tensor], Tensor]]]) – Transforms applied in the data pipeline before model input.

• augmentation_transforms (Optional[Sequence[Callable[[Tensor], Tensor]]]) – Optional augmentations applied only during training_step.

• num_classes (int) – Number of output classes.

• task_type (Optional[Literal['binary', 'multiclass', 'multilabel']]) – Override the AUROC task type; inferred from num_classes when None.

training_step(batch, batch_idx)

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

validation_step(batch, batch_idx, dataloader_idx=0)

Operates on a single batch of data from the validation set. In this step you’d might generate examples or calculate anything of interest like accuracy.

Parameters:

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

Returns:

• Tensor - The loss tensor

• dict - A dictionary. Can include any keys, but must include the key 'loss'.

• None - Skip to the next batch.

# if you have one val dataloader:
def validation_step(self, batch, batch_idx): ...


# if you have multiple val dataloaders:
def validation_step(self, batch, batch_idx, dataloader_idx=0): ...

Examples:

# CASE 1: A single validation dataset
def validation_step(self, batch, batch_idx):
    x, y = batch

    # implement your own
    out = self(x)
    loss = self.loss(out, y)

    # log 6 example images
    # or generated text... or whatever
    sample_imgs = x[:6]
    grid = torchvision.utils.make_grid(sample_imgs)
    self.logger.experiment.add_image('example_images', grid, 0)

    # calculate acc
    labels_hat = torch.argmax(out, dim=1)
    val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)

    # log the outputs!
    self.log_dict({'val_loss': loss, 'val_acc': val_acc})

If you pass in multiple val dataloaders, validation_step() will have an additional argument. We recommend setting the default value of 0 so that you can quickly switch between single and multiple dataloaders.

# CASE 2: multiple validation dataloaders
def validation_step(self, batch, batch_idx, dataloader_idx=0):
    # dataloader_idx tells you which dataset this is.
    x, y = batch

    # implement your own
    out = self(x)

    if dataloader_idx == 0:
        loss = self.loss0(out, y)
    else:
        loss = self.loss1(out, y)

    # calculate acc
    labels_hat = torch.argmax(out, dim=1)
    acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)

    # log the outputs separately for each dataloader
    self.log_dict({f"val_loss_{dataloader_idx}": loss, f"val_acc_{dataloader_idx}": acc})

Note

If you don’t need to validate you don’t need to implement this method.

Note

When the validation_step() is called, the model has been put in eval mode and PyTorch gradients have been disabled. At the end of validation, the model goes back to training mode and gradients are enabled.

predict_step(batch, batch_idx, dataloader_idx=0)

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.

Parameters:

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

Returns:

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)