# SPDX-FileCopyrightText: Copyright © 2026 Idiap Research Institute <contact@idiap.ch>
#
# SPDX-License-Identifier: GPL-3.0-or-later
"""Vision Transformer + GRU for temporal classification.
``architecture``, ``img_size``, and ``global_pool`` must match the timm backbone
and the ``Resize`` in ``model_transforms`` (see ``mednet.config.temporal.models.vitgru``).
"""
import logging
import typing
from typing import Literal
import timm
import torch
import torch.nn
import torch.optim
import torch.optim.optimizer
import torchvision.transforms
from ...data.typing import TransformSequence
from .model import Model
logger = logging.getLogger(__name__)
# Matches ``mednet.data.temporal.seq_angioreport.PHASE_ID_PADDING``.
_PHASE_ID_PADDING = -2
[docs]
def apply_backbone_unfreeze_last_n(backbone: torch.nn.Module, n: int | None) -> None:
"""Set ``requires_grad`` on a timm ViT-like ``backbone`` from its ``blocks`` only.
Parameters
----------
backbone
Module with a ``blocks`` ``ModuleList`` (timm vision transformer).
n
If ``None`` or ``n >= len(blocks)``, every parameter in ``backbone`` is
trainable. Otherwise all backbone parameters are frozen, then the last
``n`` block modules are unfrozen (``n == 0`` freezes the full backbone).
Raises
------
TypeError
If ``backbone`` has no ``blocks`` attribute.
ValueError
If ``n`` is negative or greater than ``len(blocks)``.
"""
blocks = getattr(backbone, "blocks", None)
if blocks is None:
raise TypeError(
f"Expected a timm ViT backbone with `.blocks`; got {type(backbone).__name__}."
)
depth = len(blocks)
if n is not None and n < 0:
raise ValueError(f"n must be >= 0 or None, got {n}")
if n is not None and n > depth:
raise ValueError(
f"unfreeze_last_n_backbone_blocks={n} exceeds backbone depth ({depth})."
)
if n is None or n >= depth:
for p in backbone.parameters():
p.requires_grad = True
return
for p in backbone.parameters():
p.requires_grad = False
for i in range(depth - n, depth):
for p in blocks[i].parameters():
p.requires_grad = True
[docs]
class ViTGRU(Model):
"""Vision Transformer backbone followed by a GRU temporal head.
Parameters
----------
loss_type
Loss to be used for training and evaluation.
loss_arguments
Arguments to the loss.
optimizer_type
Optimizer type for training.
optimizer_arguments
Optimizer arguments after ``params``.
scheduler_type
Optional scheduler type.
scheduler_arguments
Optional scheduler arguments after ``optimizer``.
model_transforms
Transforms to apply in the data pipeline before model input.
augmentation_transforms
Optional augmentations applied in ``training_step``.
architecture
Name of the ViT architecture to instantiate from ``timm``.
pretrained
If set to True, loads pretrained backbone weights from ``timm``.
img_size
Input image size.
global_pool
Pooling strategy for ViT features.
hidden_dim
Hidden size of the GRU.
num_layers
Number of GRU layers.
dropout
Dropout on top of the GRU output.
bidirectional
If set, uses a bidirectional GRU.
num_classes
Number of output classes.
drop_path_rate
Stochastic depth rate on the timm ViT backbone.
unfreeze_last_n_backbone_blocks
Train only the last *n* entries in ``backbone.blocks``; stem, global
norm, and earlier blocks stay frozen when ``n < len(blocks)``. ``None``
or ``n >= len(blocks)`` finetunes the entire backbone. ``n == 0`` keeps
the backbone frozen (GRU + classifier still train).
temporal_pooling
How to aggregate GRU outputs into an exam vector. ``"last"`` uses the final
hidden state (legacy). ``"attention"`` learns a softmax weight over all
timesteps (padding and ``one_per_phase`` padding slots are masked).
"""
def __init__(
self,
loss_type: type[torch.nn.Module] = torch.nn.BCEWithLogitsLoss,
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,
architecture: str = "vit_small_patch16_224.augreg_in21k",
pretrained: bool = True,
img_size: int = 224,
global_pool: Literal["", "avg", "avgmax", "max", "token", "map"] = "token",
hidden_dim: int = 256,
num_layers: int = 1,
dropout: float = 0.2,
bidirectional: bool = False,
num_classes: int = 1,
drop_path_rate: float = 0.0,
unfreeze_last_n_backbone_blocks: int | None = None,
temporal_pooling: Literal["last", "attention"] = "last",
):
super().__init__(
name=f"{architecture}-gru",
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,
)
self.architecture = architecture
self.pretrained = pretrained
self.img_size = img_size
self.global_pool = global_pool
self.hidden_dim = hidden_dim
self.num_layers = num_layers
self.dropout = dropout
self.bidirectional = bidirectional
self.drop_path_rate = drop_path_rate
if temporal_pooling not in ("last", "attention"):
raise ValueError(
f"temporal_pooling must be 'last' or 'attention', got {temporal_pooling!r}"
)
self.temporal_pooling = temporal_pooling
self.backbone = timm.create_model(
self.architecture,
img_size=(self.img_size, self.img_size),
global_pool=self.global_pool,
num_classes=0,
pretrained=self.pretrained,
drop_path_rate=self.drop_path_rate,
)
data_config = timm.data.resolve_model_data_config(self.backbone)
Model.normalizer.fset( # type: ignore[attr-defined]
self,
torchvision.transforms.Normalize(
mean=data_config["mean"],
std=data_config["std"],
),
)
self.embedding_dim = self.backbone.num_features
self.gru = torch.nn.GRU(
input_size=self.embedding_dim,
hidden_size=self.hidden_dim,
num_layers=self.num_layers,
batch_first=True,
dropout=self.dropout if self.num_layers > 1 else 0.0,
bidirectional=self.bidirectional,
)
self.dropout_layer = torch.nn.Dropout(self.dropout)
direction_factor = 2 if self.bidirectional else 1
if self.temporal_pooling == "attention":
self.temporal_attn = torch.nn.Linear(
self.hidden_dim * direction_factor,
1,
)
self.classifier = torch.nn.Linear(
self.hidden_dim * direction_factor,
self.num_classes,
)
apply_backbone_unfreeze_last_n(self.backbone, unfreeze_last_n_backbone_blocks)
# Persist constructor kwargs so ``load_from_checkpoint`` can rebuild the same
# architecture (notably ``num_classes``). Types and transform callables are
# omitted — callers restoring from config still attach ``model_transforms`` via
# the datamodule for prediction.
self.save_hyperparameters(
ignore=[
"loss_type",
"optimizer_type",
"scheduler_type",
"model_transforms",
"augmentation_transforms",
],
)
@Model.num_classes.setter # type: ignore[attr-defined]
def num_classes(self, v: int) -> None:
if self._num_classes != v:
direction_factor = 2 if getattr(self, "bidirectional", False) else 1
hidden_dim = getattr(self, "hidden_dim", 256)
self.classifier = torch.nn.Linear(hidden_dim * direction_factor, v)
self._num_classes = v
def _apply_augmentations(self, images: torch.Tensor) -> torch.Tensor:
if not self.augmentation_transforms.transforms:
return images
batch_size, seq_len, channels, height, width = images.shape
flattened = images.view(batch_size * seq_len, channels, height, width)
augmented = self.augmentation_transforms(flattened)
return augmented.view(batch_size, seq_len, channels, height, width)
@staticmethod
def _length_padding_mask(
seq_len: int,
lengths: torch.Tensor,
device: torch.device,
) -> torch.Tensor:
positions = torch.arange(seq_len, device=device).unsqueeze(0)
return positions >= lengths.to(device).unsqueeze(1)
def _run_gru(
self,
frame_embeddings: torch.Tensor,
lengths: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
packed = torch.nn.utils.rnn.pack_padded_sequence(
frame_embeddings,
lengths.cpu(),
batch_first=True,
enforce_sorted=False,
)
output_packed, hidden = self.gru(packed)
output, _ = torch.nn.utils.rnn.pad_packed_sequence(
output_packed,
batch_first=True,
)
return output, hidden
def _pool_gru_outputs(
self,
output: torch.Tensor,
hidden: torch.Tensor,
lengths: torch.Tensor,
phase_ids: torch.Tensor | None,
*,
return_attention: bool = False,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
if self.temporal_pooling == "last":
if self.bidirectional:
pooled = torch.cat((hidden[-2], hidden[-1]), dim=1)
else:
pooled = hidden[-1]
if return_attention:
raise ValueError(
"return_attention=True requires temporal_pooling='attention'"
)
return pooled
scores = self.temporal_attn(output).squeeze(-1)
mask = self._length_padding_mask(output.shape[1], lengths, output.device)
if phase_ids is not None:
mask = mask | (phase_ids.to(output.device) == _PHASE_ID_PADDING)
scores = scores.masked_fill(mask, float("-inf"))
weights = torch.softmax(scores, dim=1)
pooled = torch.sum(weights.unsqueeze(-1) * output, dim=1)
if return_attention:
return pooled, weights
return pooled
[docs]
def forward_explain(
self,
images: torch.Tensor,
lengths: torch.Tensor,
phase_ids: torch.Tensor | None = None,
*,
return_prefix_logits: bool = False,
) -> (
tuple[torch.Tensor, torch.Tensor]
| tuple[torch.Tensor, torch.Tensor, torch.Tensor]
):
"""Like :meth:`forward` but also returns attention weights ``(B, T)``.
For ``temporal_pooling == "attention"`` the weights are the learned softmax
attention over timesteps. For ``temporal_pooling == "last"`` they are
one-hot on the pooled frame (the last real timestep), so the same per-frame
reporting works for both. Padding timesteps receive weight zero.
If ``return_prefix_logits`` is True, also returns **prefix logits**
``(B, T, C)``: the classifier applied to each GRU output timestep (running
summary state after each frame). Padding positions are ``nan``.
Parameters
----------
images
Input tensor shaped as ``(batch, time, channels, height, width)``.
lengths
Real sequence lengths for each batch sample.
phase_ids
Optional per-frame phase ids ``(batch, time)``.
return_prefix_logits
If ``True``, also return per-timestep prefix logits ``(B, T, C)``.
Returns
-------
tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor, torch.Tensor]
``(logits, weights)`` or ``(logits, weights, prefix_logits)`` when
*return_prefix_logits* is ``True``.
"""
batch_size, seq_len, channels, height, width = images.shape
flat_images = images.view(batch_size * seq_len, channels, height, width)
flat_images = self.normalizer(flat_images)
frame_embeddings = self.backbone(flat_images)
frame_embeddings = frame_embeddings.view(
batch_size, seq_len, self.embedding_dim
)
if phase_ids is not None:
# Zero out ``one_per_phase`` padding slots before the GRU.
pad = phase_ids == _PHASE_ID_PADDING
frame_embeddings = frame_embeddings * (~pad).unsqueeze(-1).to(
dtype=frame_embeddings.dtype
)
output, hidden = self._run_gru(frame_embeddings, lengths)
if self.temporal_pooling == "attention":
pooled, weights = self._pool_gru_outputs(
output,
hidden,
lengths,
phase_ids,
return_attention=True,
)
else:
# 'last' pooling uses the final real timestep; mark it one-hot.
pooled = self._pool_gru_outputs(output, hidden, lengths, phase_ids)
weights = torch.zeros(
batch_size, seq_len, device=output.device, dtype=output.dtype
)
last_idx = (lengths.to(output.device) - 1).clamp(min=0, max=seq_len - 1)
weights[torch.arange(batch_size, device=output.device), last_idx] = 1.0
logits = self.classifier(self.dropout_layer(pooled))
mask = self._length_padding_mask(seq_len, lengths, weights.device)
if phase_ids is not None:
mask = mask | (phase_ids.to(weights.device) == _PHASE_ID_PADDING)
weights = weights.masked_fill(mask, 0.0)
if return_prefix_logits:
prefix_logits = self.classifier(self.dropout_layer(output))
prefix_logits = prefix_logits.masked_fill(mask.unsqueeze(-1), float("nan"))
return logits, weights, prefix_logits
return logits, weights
[docs]
def forward(
self,
images: torch.Tensor,
lengths: torch.Tensor,
phase_ids: torch.Tensor | None = None,
) -> torch.Tensor:
"""Run temporal inference.
Parameters
----------
images
Input tensor shaped as ``(batch, time, channels, height, width)``.
lengths
Real sequence lengths for each batch sample.
phase_ids
Optional per-frame phase ids ``(batch, time)``. ``one_per_phase``
padding slots (``_PHASE_ID_PADDING``) are zeroed before the GRU and
masked in attention pooling.
Returns
-------
torch.Tensor
Class logits shaped as ``(batch, num_classes)``.
"""
batch_size, seq_len, channels, height, width = images.shape
flat_images = images.view(batch_size * seq_len, channels, height, width)
flat_images = self.normalizer(flat_images)
frame_embeddings = self.backbone(flat_images)
frame_embeddings = frame_embeddings.view(
batch_size, seq_len, self.embedding_dim
)
if phase_ids is not None:
# Zero out ``one_per_phase`` padding slots before the GRU.
pad = phase_ids == _PHASE_ID_PADDING
frame_embeddings = frame_embeddings * (~pad).unsqueeze(-1).to(
dtype=frame_embeddings.dtype
)
output, hidden = self._run_gru(frame_embeddings, lengths)
pooled = self._pool_gru_outputs(output, hidden, lengths, phase_ids)
return self.classifier(self.dropout_layer(pooled)) # type: ignore[return-value]
[docs]
def training_step(self, batch, batch_idx):
del batch_idx
images = self._apply_augmentations(batch["image"])
logits = self(images, batch["lengths"], batch.get("phase_id"))
return self.train_loss(logits, batch["target"])
[docs]
def validation_step(self, batch, batch_idx, dataloader_idx=0):
del batch_idx, dataloader_idx
logits = self(batch["image"], batch["lengths"], batch.get("phase_id"))
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"], batch.get("phase_id")),
dim=1,
)