kaiko-eva 0.1.6__tar.gz → 0.1.7__tar.gz

{kaiko_eva-0.1.6 → kaiko_eva-0.1.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: kaiko-eva
-Version: 0.1.6
+Version: 0.1.7
 Summary: Evaluation Framework for oncology foundation models.
 Keywords: machine-learning,evaluation-framework,oncology,foundation-models
 Author-Email: Ioannis Gatopoulos <ioannis@kaiko.ai>, =?utf-8?q?Nicolas_K=C3=A4nzig?= <nicolas@kaiko.ai>, Roman Moser <roman@kaiko.ai>
@@ -227,7 +227,7 @@ Requires-Dist: onnxruntime>=1.15.1
 Requires-Dist: onnx>=1.16.0
 Requires-Dist: toolz>=0.12.1
 Requires-Dist: rich>=13.7.1
-Requires-Dist: torchmetrics>=1.4.2
+Requires-Dist: torchmetrics>=1.6.0
 Provides-Extra: vision
 Requires-Dist: h5py>=3.10.0; extra == "vision"
 Requires-Dist: nibabel>=4.0.1; extra == "vision"

{kaiko_eva-0.1.6 → kaiko_eva-0.1.7}/pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "pdm.backend"
 
 [project]
 name = "kaiko-eva"
-version = "0.1.6"
+version = "0.1.7"
 description = "Evaluation Framework for oncology foundation models."
 keywords = [
     "machine-learning",
@@ -45,7 +45,7 @@ dependencies = [
     "onnx>=1.16.0",
     "toolz>=0.12.1",
     "rich>=13.7.1",
-    "torchmetrics>=1.4.2",
+    "torchmetrics>=1.6.0",
 ]
 
 [project.urls]

{kaiko_eva-0.1.6 → kaiko_eva-0.1.7}/src/eva/core/data/dataloaders/dataloader.py

@@ -59,17 +59,20 @@ class DataLoader:
     prefetch_factor: int | None = 2
     """Number of batches loaded in advance by each worker."""
 
-    def __call__(self, dataset: datasets.TorchDataset) -> dataloader.DataLoader:
+    def __call__(
+        self, dataset: datasets.TorchDataset, sampler: samplers.Sampler | None = None
+    ) -> dataloader.DataLoader:
         """Returns the dataloader on the provided dataset.
 
         Args:
             dataset: dataset from which to load the data.
+            sampler: defines the strategy to draw samples from the dataset.
         """
         return dataloader.DataLoader(
             dataset=dataset,
             batch_size=self.batch_size,
             shuffle=self.shuffle,
-            sampler=self.sampler,
+            sampler=sampler or self.sampler,
             batch_sampler=self.batch_sampler,
             num_workers=self.num_workers or multiprocessing.cpu_count(),
             collate_fn=self.collate_fn,

{kaiko_eva-0.1.6 → kaiko_eva-0.1.7}/src/eva/core/data/datamodules/datamodule.py

@@ -8,6 +8,7 @@ from typing_extensions import override
 
 from eva.core.data import dataloaders as dataloaders_lib
 from eva.core.data import datasets as datasets_lib
+from eva.core.data import samplers as samplers_lib
 from eva.core.data.datamodules import call, schemas
 
 
@@ -24,17 +25,20 @@ class DataModule(pl.LightningDataModule):
         self,
         datasets: schemas.DatasetsSchema | None = None,
         dataloaders: schemas.DataloadersSchema | None = None,
+        samplers: schemas.SamplersSchema | None = None,
     ) -> None:
         """Initializes the datamodule.
 
         Args:
             datasets: The desired datasets.
             dataloaders: The desired dataloaders.
+            samplers: The desired samplers for the dataloaders.
         """
         super().__init__()
 
         self.datasets = datasets or self.default_datasets
         self.dataloaders = dataloaders or self.default_dataloaders
+        self.samplers = samplers or self.default_samplers
 
     @property
     def default_datasets(self) -> schemas.DatasetsSchema:
@@ -46,6 +50,11 @@ class DataModule(pl.LightningDataModule):
         """Returns the default dataloader schema."""
         return schemas.DataloadersSchema()
 
+    @property
+    def default_samplers(self) -> schemas.SamplersSchema:
+        """Returns the default samplers schema."""
+        return schemas.SamplersSchema()
+
     @override
     def prepare_data(self) -> None:
         call.call_method_if_exists(self.datasets.tolist(), "prepare_data")
@@ -64,7 +73,12 @@ class DataModule(pl.LightningDataModule):
             raise ValueError(
                 "Train dataloader can not be initialized as `self.datasets.train` is `None`."
             )
-        return self.dataloaders.train(self.datasets.train)
+        if isinstance(self.datasets.train, list) and len(self.datasets.train) > 1:
+            raise ValueError("Train dataloader can not be initialized with multiple datasets.")
+
+        return self._initialize_dataloaders(
+            self.dataloaders.train, self.datasets.train, self.samplers.train
+        )[0]
 
     @override
     def val_dataloader(self) -> EVAL_DATALOADERS:
@@ -72,7 +86,9 @@ class DataModule(pl.LightningDataModule):
             raise ValueError(
                 "Validation dataloader can not be initialized as `self.datasets.val` is `None`."
             )
-        return self._initialize_dataloaders(self.dataloaders.val, self.datasets.val)
+        return self._initialize_dataloaders(
+            self.dataloaders.val, self.datasets.val, self.samplers.val
+        )
 
     @override
     def test_dataloader(self) -> EVAL_DATALOADERS:
@@ -80,7 +96,9 @@ class DataModule(pl.LightningDataModule):
             raise ValueError(
                 "Test dataloader can not be initialized as `self.datasets.test` is `None`."
             )
-        return self._initialize_dataloaders(self.dataloaders.test, self.datasets.test)
+        return self._initialize_dataloaders(
+            self.dataloaders.test, self.datasets.test, self.samplers.test
+        )
 
     @override
     def predict_dataloader(self) -> EVAL_DATALOADERS:
@@ -88,21 +106,40 @@ class DataModule(pl.LightningDataModule):
             raise ValueError(
                 "Predict dataloader can not be initialized as `self.datasets.predict` is `None`."
             )
-        return self._initialize_dataloaders(self.dataloaders.predict, self.datasets.predict)
+        if isinstance(self.datasets.predict, list) and len(self.datasets.predict) > 1:
+            # Only apply sampler to the first predict dataset (should correspond to train split)
+            train_dataloader = self._initialize_dataloaders(
+                self.dataloaders.predict, self.datasets.predict[0], self.samplers.predict
+            )
+            return train_dataloader + self._initialize_dataloaders(
+                self.dataloaders.predict, self.datasets.predict[1:]
+            )
+
+        return self._initialize_dataloaders(
+            self.dataloaders.predict, self.datasets.predict, self.samplers.predict
+        )
 
     def _initialize_dataloaders(
         self,
         dataloader: dataloaders_lib.DataLoader,
         datasets: datasets_lib.TorchDataset | List[datasets_lib.TorchDataset],
+        sampler: samplers_lib.Sampler | None = None,
     ) -> EVAL_DATALOADERS:
         """Initializes dataloaders from a given set of dataset.
 
         Args:
             dataloader: The dataloader to apply to the provided datasets.
             datasets: The desired dataset(s) to allocate dataloader(s).
+            sampler: The sampler to use for the dataloader.
 
         Returns:
             A list with the dataloaders of the provided dataset(s).
         """
         datasets = datasets if isinstance(datasets, list) else [datasets]
-        return list(map(dataloader, datasets))
+
+        dataloaders = []
+        for dataset in datasets:
+            if sampler is not None and isinstance(sampler, samplers_lib.SamplerWithDataSource):
+                sampler.set_dataset(dataset)  # type: ignore
+            dataloaders.append(dataloader(dataset, sampler=sampler))
+        return dataloaders

{kaiko_eva-0.1.6 → kaiko_eva-0.1.7}/src/eva/core/data/datamodules/schemas.py

@@ -3,7 +3,7 @@
 import dataclasses
 from typing import List
 
-from eva.core.data import dataloaders, datasets
+from eva.core.data import dataloaders, datasets, samplers
 
 TRAIN_DATASET = datasets.TorchDataset | None
 """Train dataset."""
@@ -60,3 +60,20 @@ class DataloadersSchema:
 
     predict: dataloaders.DataLoader = dataclasses.field(default_factory=dataloaders.DataLoader)
     """Predict dataloader."""
+
+
+@dataclasses.dataclass(frozen=True)
+class SamplersSchema:
+    """Samplers schema used in DataModule."""
+
+    train: samplers.Sampler | None = None
+    """Train sampler."""
+
+    val: samplers.Sampler | None = None
+    """Validation sampler."""
+
+    test: samplers.Sampler | None = None
+    """Test sampler."""
+
+    predict: samplers.Sampler | None = None
+    """Predict sampler."""

{kaiko_eva-0.1.6 → kaiko_eva-0.1.7}/src/eva/core/data/datasets/__init__.py

@@ -1,15 +1,18 @@
 """Datasets API."""
 
-from eva.core.data.datasets.base import Dataset
+from eva.core.data.datasets.base import Dataset, MapDataset
 from eva.core.data.datasets.classification import (
     EmbeddingsClassificationDataset,
     MultiEmbeddingsClassificationDataset,
 )
 from eva.core.data.datasets.dataset import TorchDataset
+from eva.core.data.datasets.typings import DataSample
 
 __all__ = [
     "Dataset",
+    "MapDataset",
     "EmbeddingsClassificationDataset",
     "MultiEmbeddingsClassificationDataset",
     "TorchDataset",
+    "DataSample",
 ]

{kaiko_eva-0.1.6 → kaiko_eva-0.1.7}/src/eva/core/data/datasets/base.py

@@ -1,5 +1,7 @@
 """Base dataset class."""
 
+import abc
+
 from eva.core.data.datasets import dataset
 
 
@@ -51,3 +53,24 @@ class Dataset(dataset.TorchDataset):
         of fit (train + validate), validate, test, or predict and it will be
         called from every process (i.e. GPU) across all the nodes in DDP.
         """
+
+
+class MapDataset(Dataset):
+    """Abstract base class for all map-style datasets."""
+
+    @abc.abstractmethod
+    def __getitem__(self, index: int):
+        """Retrieves the item at the given index.
+
+        Args:
+            index: Index of the item to retrieve.
+
+        Returns:
+            The data at the given index.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def __len__(self) -> int:
+        """Returns the length of the dataset."""
+        raise NotImplementedError

kaiko_eva-0.1.7/src/eva/core/data/datasets/typings.py

@@ -0,0 +1,18 @@
+"""Typing definitions for the datasets module."""
+
+from typing import Any, Dict, NamedTuple
+
+import torch
+
+
+class DataSample(NamedTuple):
+    """The default input batch data scheme."""
+
+    data: torch.Tensor
+    """The data batch."""
+
+    targets: torch.Tensor | None = None
+    """The target batch."""
+
+    metadata: Dict[str, Any] | None = None
+    """The associated metadata."""

kaiko_eva-0.1.7/src/eva/core/data/samplers/__init__.py

@@ -0,0 +1,7 @@
+"""Data samplers API."""
+
+from eva.core.data.samplers.classification.balanced import BalancedSampler
+from eva.core.data.samplers.random import RandomSampler
+from eva.core.data.samplers.sampler import Sampler, SamplerWithDataSource
+
+__all__ = ["Sampler", "SamplerWithDataSource", "RandomSampler", "BalancedSampler"]

kaiko_eva-0.1.7/src/eva/core/data/samplers/classification/__init__.py

@@ -0,0 +1,5 @@
+"""Classification data samplers API."""
+
+from eva.core.data.samplers.classification.balanced import BalancedSampler
+
+__all__ = ["BalancedSampler"]

kaiko_eva-0.1.7/src/eva/core/data/samplers/classification/balanced.py

@@ -0,0 +1,96 @@
+"""Random class sampler for data loading."""
+
+from collections import defaultdict
+from typing import Dict, Iterator, List
+
+import numpy as np
+from typing_extensions import override
+
+from eva.core.data import datasets
+from eva.core.data.datasets.typings import DataSample
+from eva.core.data.samplers.sampler import SamplerWithDataSource
+from eva.core.utils.progress_bar import tqdm
+
+
+class BalancedSampler(SamplerWithDataSource[int]):
+    """Balanced class sampler for data loading.
+
+    The sampler ensures that:
+    1. Each class has the same number of samples
+    2. Samples within each class are randomly selected
+    3. Samples of different classes appear in random order
+    """
+
+    def __init__(self, num_samples: int, replacement: bool = False, seed: int | None = 42):
+        """Initializes the balanced sampler.
+
+        Args:
+            num_samples: The number of samples to draw per class.
+            replacement: samples are drawn on-demand with replacement if ``True``, default=``False``
+            seed: Random seed for reproducibility.
+        """
+        self._num_samples = num_samples
+        self._replacement = replacement
+        self._class_indices: Dict[int, List[int]] = defaultdict(list)
+        self._random_generator = np.random.default_rng(seed)
+
+    def __len__(self) -> int:
+        """Returns the total number of samples."""
+        return self._num_samples * len(self._class_indices)
+
+    def __iter__(self) -> Iterator[int]:
+        """Creates an iterator that yields indices in a class balanced way.
+
+        Returns:
+            Iterator yielding dataset indices.
+        """
+        indices = []
+
+        for class_idx in self._class_indices:
+            class_indices = self._class_indices[class_idx]
+            sampled_indices = self._random_generator.choice(
+                class_indices, size=self._num_samples, replace=self._replacement
+            ).tolist()
+            indices.extend(sampled_indices)
+
+        self._random_generator.shuffle(indices)
+
+        return iter(indices)
+
+    @override
+    def set_dataset(self, data_source: datasets.MapDataset):
+        """Sets the dataset and builds class indices.
+
+        Args:
+            data_source: The dataset to sample from.
+
+        Raises:
+            ValueError: If the dataset doesn't have targets or if any class has
+                fewer samples than `num_samples` and `replacement` is `False`.
+        """
+        super().set_dataset(data_source)
+        self._make_indices()
+
+    def _make_indices(self):
+        """Builds indices for each class in the dataset."""
+        self._class_indices.clear()
+
+        for idx in tqdm(
+            range(len(self.data_source)), desc="Fetching class indices for balanced sampler"
+        ):
+            _, target, _ = DataSample(*self.data_source[idx])
+            if target is None:
+                raise ValueError("The dataset must return non-empty targets.")
+            if target.numel() != 1:
+                raise ValueError("The dataset must return a single & scalar target.")
+
+            class_idx = int(target.item())
+            self._class_indices[class_idx].append(idx)
+
+        if not self._replacement:
+            for class_idx, indices in self._class_indices.items():
+                if len(indices) < self._num_samples:
+                    raise ValueError(
+                        f"Class {class_idx} has only {len(indices)} samples, "
+                        f"which is less than the required {self._num_samples} samples."
+                    )

kaiko_eva-0.1.7/src/eva/core/data/samplers/random.py

@@ -0,0 +1,39 @@
+"""Random sampler for data loading."""
+
+from typing import Optional
+
+from torch.utils import data
+from typing_extensions import override
+
+from eva.core.data import datasets
+from eva.core.data.samplers.sampler import SamplerWithDataSource
+
+
+class RandomSampler(data.RandomSampler, SamplerWithDataSource[int]):
+    """Samples elements randomly."""
+
+    data_source: datasets.MapDataset  # type: ignore
+
+    def __init__(
+        self, replacement: bool = False, num_samples: Optional[int] = None, generator=None
+    ) -> None:
+        """Initializes the random sampler.
+
+        Args:
+            data_source: dataset to sample from
+            replacement: samples are drawn on-demand with replacement if ``True``, default=``False``
+            num_samples: number of samples to draw, default=`len(dataset)`.
+            generator: Generator used in sampling.
+        """
+        self.replacement = replacement
+        self._num_samples = num_samples
+        self.generator = generator
+
+    @override
+    def set_dataset(self, data_source: datasets.MapDataset) -> None:
+        super().__init__(
+            data_source,
+            replacement=self.replacement,
+            num_samples=self.num_samples,
+            generator=self.generator,
+        )

kaiko_eva-0.1.7/src/eva/core/data/samplers/sampler.py

@@ -0,0 +1,33 @@
+"""Core data sampler."""
+
+from typing import Generic, TypeVar
+
+from torch.utils import data
+
+from eva.core.data import datasets
+
+Sampler = data.Sampler
+"""Core abstract data sampler class."""
+
+T_co = TypeVar("T_co", covariant=True)
+
+
+class SamplerWithDataSource(Sampler, Generic[T_co]):
+    """A sampler base class that enables to specify the data source after initialization.
+
+    The `set_dataset` can also be overwritten to expand the functionality of the derived
+    sampler classes.
+    """
+
+    data_source: datasets.MapDataset
+
+    def set_dataset(self, data_source: datasets.MapDataset) -> None:
+        """Sets the dataset to sample from.
+
+        This is not done in the constructor because the dataset might not be
+        available at that time.
+
+        Args:
+            data_source: The dataset to sample from.
+        """
+        self.data_source = data_source

{kaiko_eva-0.1.6 → kaiko_eva-0.1.7}/src/eva/core/metrics/structs/module.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+from typing import List
+
 from torch import nn
 
 from eva.core.metrics.structs import collection, schemas
@@ -46,6 +48,7 @@ class MetricModule(nn.Module):
         test: MetricModuleType | None,
         *,
         separator: str = "/",
+        compute_groups: bool | List[List[str]] = True,
     ) -> MetricModule:
         """Initializes a metric module from a list of metrics.
 
@@ -55,19 +58,25 @@ class MetricModule(nn.Module):
             test: Metrics for the test stage.
             separator: The separator between the group name of the metric
                 and the metric itself.
+            compute_groups: All metrics in a compute group share the same metric state
+                and are therefore only different in their compute step. To disable this
+                behavior, set to `False`.
         """
         return cls(
-            train=_create_collection_from_metrics(train, prefix="train" + separator),
-            val=_create_collection_from_metrics(val, prefix="val" + separator),
-            test=_create_collection_from_metrics(test, prefix="test" + separator),
+            train=_create_collection_from_metrics(
+                train, prefix="train" + separator, compute_groups=compute_groups
+            ),
+            val=_create_collection_from_metrics(
+                val, prefix="val" + separator, compute_groups=compute_groups
+            ),
+            test=_create_collection_from_metrics(
+                test, prefix="test" + separator, compute_groups=compute_groups
+            ),
         )
 
     @classmethod
     def from_schema(
-        cls,
-        schema: schemas.MetricsSchema,
-        *,
-        separator: str = "/",
+        cls, schema: schemas.MetricsSchema, *, separator: str = "/", compute_groups: bool = True
     ) -> MetricModule:
         """Initializes a metric module from the metrics schema.
 
@@ -75,12 +84,16 @@ class MetricModule(nn.Module):
             schema: The dataclass metric schema.
             separator: The separator between the group name of the metric
                 and the metric itself.
+            compute_groups: All metrics in a compute group share the same metric state
+                and are therefore only different in their compute step. To disable this
+                behavior, set to `False`.
         """
         return cls.from_metrics(
             train=schema.training_metrics,
             val=schema.evaluation_metrics,
             test=schema.evaluation_metrics,
             separator=separator,
+            compute_groups=compute_groups,
         )
 
     @property
@@ -100,16 +113,24 @@ class MetricModule(nn.Module):
 
 
 def _create_collection_from_metrics(
-    metrics: MetricModuleType | None, *, prefix: str | None = None
+    metrics: MetricModuleType | None,
+    *,
+    prefix: str | None = None,
+    compute_groups: bool | List[List[str]] = True,
 ) -> collection.MetricCollection | None:
     """Create a unique collection from metrics.
 
     Args:
         metrics: The desired metrics.
         prefix: A prefix to added to the collection.
+        compute_groups: All metrics in a compute group share the same metric state
+            and are therefore only different in their compute step. To disable this
+            behavior, set to `False`.
 
     Returns:
         The resulted metrics collection.
     """
-    metrics_collection = collection.MetricCollection(metrics or [], prefix=prefix)  # type: ignore
+    metrics_collection = collection.MetricCollection(
+        metrics=metrics or [], prefix=prefix, compute_groups=compute_groups  # type: ignore
+    )
     return metrics_collection.clone()

{kaiko_eva-0.1.6 → kaiko_eva-0.1.7}/src/eva/core/models/__init__.py

@@ -2,7 +2,13 @@
 
 from eva.core.models.modules import HeadModule, InferenceModule
 from eva.core.models.networks import MLP
-from eva.core.models.wrappers import BaseModel, HuggingFaceModel, ModelFromFunction, ONNXModel
+from eva.core.models.wrappers import (
+    BaseModel,
+    HuggingFaceModel,
+    ModelFromFunction,
+    ONNXModel,
+    TorchHubModel,
+)
 
 __all__ = [
     "HeadModule",
@@ -12,4 +18,5 @@ __all__ = [
     "HuggingFaceModel",
     "ModelFromFunction",
     "ONNXModel",
+    "TorchHubModel",
 ]

{kaiko_eva-0.1.6 → kaiko_eva-0.1.7}/src/eva/core/models/modules/head.py

@@ -12,7 +12,7 @@ from typing_extensions import override
 from eva.core.metrics import structs as metrics_lib
 from eva.core.models.modules import module
 from eva.core.models.modules.typings import INPUT_BATCH, MODEL_TYPE
-from eva.core.models.modules.utils import batch_postprocess, grad
+from eva.core.models.modules.utils import batch_postprocess, grad, submodule_state_dict
 from eva.core.utils import parser
 
 
@@ -32,6 +32,7 @@ class HeadModule(module.ModelModule):
         lr_scheduler: LRSchedulerCallable = lr_scheduler.ConstantLR,
         metrics: metrics_lib.MetricsSchema | None = None,
         postprocess: batch_postprocess.BatchPostProcess | None = None,
+        save_head_only: bool = True,
     ) -> None:
         """Initializes the neural net head module.
 
@@ -48,6 +49,8 @@ class HeadModule(module.ModelModule):
             postprocess: A list of helper functions to apply after the
                 loss and before the metrics calculation to the model
                 predictions and targets.
+            save_head_only: Whether to save only the head during checkpointing. If False,
+                will also save the backbone (not recommended when frozen).
         """
         super().__init__(metrics=metrics, postprocess=postprocess)
 
@@ -56,6 +59,7 @@ class HeadModule(module.ModelModule):
         self.backbone = backbone
         self.optimizer = optimizer
         self.lr_scheduler = lr_scheduler
+        self.save_head_only = save_head_only
 
     @override
     def configure_model(self) -> Any:
@@ -72,6 +76,20 @@ class HeadModule(module.ModelModule):
         lr_scheduler = self.lr_scheduler(optimizer)
         return {"optimizer": optimizer, "lr_scheduler": lr_scheduler}
 
+    @override
+    def on_save_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
+        if self.save_head_only:
+            checkpoint["state_dict"] = submodule_state_dict(checkpoint["state_dict"], "head")
+        super().on_save_checkpoint(checkpoint)
+
+    @override
+    def on_load_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
+        if self.save_head_only and self.backbone is not None:
+            checkpoint["state_dict"].update(
+                {f"backbone.{k}": v for k, v in self.backbone.state_dict().items()}
+            )
+        super().on_load_checkpoint(checkpoint)
+
     @override
     def forward(self, tensor: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor:
         features = tensor if self.backbone is None else self.backbone(tensor)

{kaiko_eva-0.1.6 → kaiko_eva-0.1.7}/src/eva/core/models/modules/utils/__init__.py

@@ -2,5 +2,6 @@
 
 from eva.core.models.modules.utils import grad
 from eva.core.models.modules.utils.batch_postprocess import BatchPostProcess
+from eva.core.models.modules.utils.checkpoint import submodule_state_dict
 
-__all__ = ["grad", "BatchPostProcess"]
+__all__ = ["grad", "BatchPostProcess", "submodule_state_dict"]

kaiko_eva-0.1.7/src/eva/core/models/modules/utils/checkpoint.py

@@ -0,0 +1,21 @@
+"""Checkpointing related utilities and helper functions."""
+
+from typing import Any, Dict
+
+
+def submodule_state_dict(state_dict: Dict[str, Any], submodule_key: str) -> Dict[str, Any]:
+    """Get the state dict of a specific submodule.
+
+    Args:
+        state_dict: The state dict to extract the submodule from.
+        submodule_key: The key of the submodule to extract.
+
+    Returns:
+        The subset of the state dict corresponding to the specified submodule.
+    """
+    submodule_key = submodule_key if submodule_key.endswith(".") else submodule_key + "."
+    return {
+        module: weights
+        for module, weights in state_dict.items()
+        if module.startswith(submodule_key)
+    }