Source code for mednet.config.temporal.models.vitgru

# SPDX-FileCopyrightText: Copyright © 2026 Idiap Research Institute <contact@idiap.ch>
#
# SPDX-License-Identifier: GPL-3.0-or-later
"""ViT + GRU for sequence classification @ 224 (RetFound-Green backbone).

Hyperparameters match ``fm-overspecialization`` **RETFound-Green** +
**HyperF_Type** (``config/hyperparameters.py``): ``lr=9.68e-6``,
``weight_decay=5e-3``, no LR scheduler, ``drop_path_rate=0.1`` on the ViT.

Preprocessing matches that repo's ``configuration.set_model`` for the ViT path:
**Resize(224) + RGB** (no square pad). ``seq_finetune.experiment`` default
``--batch-size`` is **16** for the same row.

``vit_small_patch14_reg4_dinov2`` matches patch weights in
``retfoundgreen_statedict.pth``; optional ``--pretrained-weights`` uses the same
load path as ``fm-overspecialization`` (including ``pos_embed`` interpolation when
the checkpoint grid differs from 224).

``vit_small_patch14_reg4_dinov2`` has **12** transformer blocks. Use
``build_vitgru_model(unfreeze_last_n_backbone_blocks=n)`` with ``n == 12`` for full
backbone finetuning (default), or a smaller ``n`` to train only the last blocks.
"""

import torch.nn
import torch.optim
import torchvision.transforms
import torchvision.transforms.v2

import mednet.models.temporal.vitgru

_IMG = 224
# fm-overspecialization/config/hyperparameters.py — RETFound-Green, HyperF_Type
_LR = 9.68e-6
_WEIGHT_DECAY = 5e-3
_DROP_PATH = 0.1
# Depth of ``vit_small_patch14_reg4_dinov2`` ``.blocks`` (full backbone FT).
_DEFAULT_UNFREEZE_LAST_N = 12




def build_vitgru_model(
    unfreeze_last_n_backbone_blocks: int | None = None,
) -> mednet.models.temporal.vitgru.ViTGRU:
    """Instantiate the temporal ViT+GRU used by ``seq_finetune``.

    Parameters
    ----------
    unfreeze_last_n_backbone_blocks
        Forwarded to :class:`~mednet.models.temporal.vitgru.ViTGRU`. ``None`` uses
        ``_DEFAULT_UNFREEZE_LAST_N`` (12 = finetune the full RETFound-Green ViT).

    Returns
    -------
    mednet.models.temporal.vitgru.ViTGRU
        Configured :class:`~mednet.models.temporal.vitgru.ViTGRU` instance.
    """
    n = (
        _DEFAULT_UNFREEZE_LAST_N
        if unfreeze_last_n_backbone_blocks is None
        else unfreeze_last_n_backbone_blocks
    )
    return mednet.models.temporal.vitgru.ViTGRU(
        architecture="vit_small_patch14_reg4_dinov2",
        pretrained=False,
        img_size=_IMG,
        global_pool="avg",
        drop_path_rate=_DROP_PATH,
        hidden_dim=256,
        num_layers=1,
        dropout=0.2,
        bidirectional=False,
        num_classes=5,
        loss_type=torch.nn.CrossEntropyLoss,
        optimizer_type=torch.optim.AdamW,
        optimizer_arguments=dict(lr=_LR, weight_decay=_WEIGHT_DECAY),
        model_transforms=[
            torchvision.transforms.v2.Resize(
                (_IMG, _IMG),
                antialias=True,
                interpolation=torchvision.transforms.InterpolationMode.BICUBIC,
            ),
            torchvision.transforms.v2.RGB(),
        ],
        unfreeze_last_n_backbone_blocks=n,
    )



model = build_vitgru_model()