mednet.config.temporal.models.vitgru¶
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.
Functions
|
Instantiate the temporal ViT+GRU used by |
- mednet.config.temporal.models.vitgru.build_vitgru_model(unfreeze_last_n_backbone_blocks=None)[source]¶
Instantiate the temporal ViT+GRU used by
seq_finetune.
# 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()