# SPDX-FileCopyrightText: Copyright © 2026 Idiap Research Institute <contact@idiap.ch>
#
# SPDX-License-Identifier: GPL-3.0-or-later
"""Definition of base model type for temporal classification tasks."""
import logging
import typing
from typing import Literal
import torch
import torch.nn
import torch.optim
import torch.optim.optimizer
import torchmetrics.classification
from ...data.typing import TransformSequence
from ..classify.model import Model as ClassificationModel
logger = logging.getLogger(__name__)
[docs]
class Model(ClassificationModel):
"""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 :class:`mednet.models.classify.vit.ViT` extensions).
Parameters
----------
name
Human-readable model name used for logging.
loss_type
Loss module type; defaults to :class:`torch.nn.BCEWithLogitsLoss`.
loss_arguments
Keyword arguments forwarded to *loss_type*.
optimizer_type
Optimizer class; defaults to :class:`torch.optim.Adam`.
optimizer_arguments
Keyword arguments forwarded to *optimizer_type* (after ``params``).
scheduler_type
Optional LR scheduler class.
scheduler_arguments
Keyword arguments forwarded to *scheduler_type* (after ``optimizer``).
model_transforms
Transforms applied in the data pipeline before model input.
augmentation_transforms
Optional augmentations applied only during ``training_step``.
num_classes
Number of output classes.
task_type
Override the AUROC task type; inferred from *num_classes* when ``None``.
"""
def __init__(
self,
name: str,
loss_type: type[torch.nn.Module] | None = None,
loss_arguments: dict[str, typing.Any] | None = None,
optimizer_type: type[torch.optim.Optimizer] = torch.optim.Adam,
optimizer_arguments: dict[str, typing.Any] | None = None,
scheduler_type: type[torch.optim.lr_scheduler.LRScheduler] | None = None,
scheduler_arguments: dict[str, typing.Any] | None = None,
model_transforms: TransformSequence | None = None,
augmentation_transforms: TransformSequence | None = None,
num_classes: int = 1,
task_type: Literal["binary", "multiclass", "multilabel"] | None = None,
):
super().__init__(
name=name,
loss_type=loss_type,
loss_arguments=loss_arguments,
optimizer_type=optimizer_type,
optimizer_arguments=optimizer_arguments,
scheduler_type=scheduler_type,
scheduler_arguments=scheduler_arguments,
model_transforms=model_transforms,
augmentation_transforms=augmentation_transforms,
num_classes=num_classes,
)
if task_type is None:
self._task_type: Literal["binary", "multiclass", "multilabel"] = (
"binary" if num_classes == 1 else "multiclass"
)
else:
self._task_type = task_type
average = "none" if self.num_classes == 1 else "macro"
self.val_auc = torchmetrics.classification.AUROC(
task=self._task_type,
num_classes=self.num_classes,
num_labels=self.num_classes,
average=average,
)
[docs]
def training_step(self, batch, batch_idx):
del batch_idx
return self.train_loss(
self(batch["image"], batch["lengths"]),
batch["target"],
)
[docs]
def validation_step(self, batch, batch_idx, dataloader_idx=0):
del batch_idx, dataloader_idx
logits = self(batch["image"], batch["lengths"])
targets = batch["target"]
self.val_auc(logits, targets.int())
self.log(
"auc/validation",
self.val_auc,
on_step=False,
on_epoch=True,
)
return self.validation_loss(logits, targets)
[docs]
def predict_step(self, batch, batch_idx, dataloader_idx=0):
del batch_idx, dataloader_idx
return torch.softmax(self(batch["image"], batch["lengths"]), dim=1)