Use real type annotations in the hydra config

conf/experiment/nicopp/rn18/only_pred_y_loss.yaml

@@ -1,7 +1,7 @@
 # @package _global_
 
 defaults:
-    - /ae: cosine_annealing
+    - /ae_opt: cosine_annealing
     - /alg: only_pred_y_loss
     - /eval: nicopp
     - override /ae_arch: resnet

pyproject.toml

@@ -24,7 +24,7 @@ numpy = { version = ">=1.23.2" }
 pandas = { version = ">=1.5.0" }
 pillow = "*"
 python = ">=3.10,<3.13"
-ranzen = { version = "^2.1.2" }
+ranzen = { version = "^2.5.0" }
 scikit-image = ">=0.14"
 scikit_learn = { version = ">=0.20.1" }
 scipy = { version = ">=1.2.1" }

src/algs/adv/scorer.py

@@ -66,6 +66,7 @@ def balanced_accuracy(y_pred: Tensor, *, y_true: Tensor) -> Tensor:
     return cdtm.subclass_balanced_accuracy(y_pred=y_pred, y_true=y_true, s=y_true)
 
 
+@dataclass(eq=False)
 class Scorer(ABC):
     @abstractmethod
     def run(

src/arch/predictors/base.py

@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+from dataclasses import dataclass
 from typing import TypeVar
 from typing_extensions import TypeAliasType
 
@@ -10,6 +11,7 @@
 PredictorFactoryOut = TypeAliasType("PredictorFactoryOut", tuple[M, int], type_params=(M,))
 
 
+@dataclass(eq=False)
 class PredictorFactory(ABC):
     @abstractmethod
     def __call__(

src/data/common.py

@@ -75,6 +75,7 @@ def num_samples_te(self) -> int:
         return len(self.test)
 
 
+@dataclass
 class DatasetFactory(ABC):
     @abstractmethod
     def __call__(self) -> Dataset:

src/labelling/pipeline.py

@@ -38,6 +38,7 @@
 ]
 
 
+@dataclass(repr=False, eq=False)
 class Labeller(ABC):
     @abstractmethod
     def run(self, dm: DataModule, device: torch.device) -> Tensor | None:

src/relay/base.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass, field
 import os
-from typing import Any, ClassVar
+from typing import ClassVar
 
 from loguru import logger
 from ranzen.torch import random_seed
@@ -18,7 +18,7 @@
 @dataclass(eq=False, kw_only=True)
 class BaseRelay:
     dm: DataModuleConf = field(default_factory=DataModuleConf)
-    split: Any
+    split: DataSplitter
     wandb: WandbConf = field(default_factory=WandbConf)
     seed: int = 0
 
@@ -36,8 +36,6 @@ def init_dm(
         labeller: Labeller,
         device: torch.device,
     ) -> DataModule:
-        assert isinstance(self.split, DataSplitter)
-
         logger.info(f"Current working directory: '{os.getcwd()}'")
         random_seed(self.seed, use_cuda=True)
         torch.multiprocessing.set_sharing_strategy("file_system")

src/relay/fs.py

@@ -37,11 +37,11 @@ class FsRelay(BaseRelay):
         ]
     )
 
-    alg: Any
-    ds: Any
-    backbone: Any
+    alg: FsAlg
+    ds: DatasetFactory
+    backbone: BackboneFactory
     predictor: Fcn = field(default_factory=Fcn)
-    labeller: Any
+    labeller: Labeller
 
     options: ClassVar[dict[str, dict[str, type]]] = BaseRelay.options | {
         "ds": {
@@ -70,11 +70,6 @@ class FsRelay(BaseRelay):
     }
 
     def run(self, raw_config: dict[str, Any] | None = None) -> float | None:
-        assert isinstance(self.alg, FsAlg)
-        assert isinstance(self.backbone, BackboneFactory)
-        assert isinstance(self.ds, DatasetFactory)
-        assert isinstance(self.labeller, Labeller)
-
         ds = self.ds()
         run = self.wandb.init(raw_config, (ds, self.labeller, self.backbone, self.predictor))
         dm = self.init_dm(ds, self.labeller, device=self.alg.device)

src/relay/label.py

@@ -29,8 +29,8 @@ class LabelRelay(BaseRelay):
         ]
     )
 
-    ds: Any  # CdtDataset
-    labeller: Any  # Labeller
+    ds: DatasetFactory
+    labeller: Labeller
     gpu: int = 0
 
     options: ClassVar[dict[str, dict[str, type]]] = BaseRelay.options | {
@@ -50,9 +50,6 @@ class LabelRelay(BaseRelay):
     }
 
     def run(self, raw_config: dict[str, Any] | None = None) -> float | None:
-        assert isinstance(self.ds, DatasetFactory)
-        assert isinstance(self.labeller, Labeller)
-
         ds = self.ds()
         run = self.wandb.init(raw_config, (ds, self.labeller))
         device = resolve_device(self.gpu)

src/relay/mimin.py

@@ -37,14 +37,14 @@ class MiMinRelay(BaseRelay):
     )
 
     alg: MiMin = field(default_factory=MiMin)
-    ae_arch: Any
+    ae_arch: AeFactory
     disc_arch: Fcn = field(default_factory=Fcn)
     disc: OptimizerCfg = field(default_factory=OptimizerCfg)
     eval: Evaluator = field(default_factory=Evaluator)
     ae: SplitAeCfg = field(default_factory=SplitAeCfg)
     ae_opt: OptimizerCfg = field(default_factory=OptimizerCfg)
-    ds: Any
-    labeller: Any
+    ds: DatasetFactory
+    labeller: Labeller
 
     options: ClassVar[dict[str, dict[str, type]]] = BaseRelay.options | {
         "ds": {
@@ -70,10 +70,6 @@ class MiMinRelay(BaseRelay):
     }
 
     def run(self, raw_config: dict[str, Any] | None = None) -> None:
-        assert isinstance(self.ae_arch, AeFactory)
-        assert isinstance(self.ds, DatasetFactory)
-        assert isinstance(self.labeller, Labeller)
-
         ds = self.ds()
         run = self.wandb.init(raw_config, (ds, self.labeller, self.ae_arch, self.disc_arch))
         dm = self.init_dm(ds, self.labeller, device=self.alg.device)

src/relay/split.py

@@ -1,7 +1,7 @@
 from dataclasses import dataclass, field
 from typing import Any, ClassVar
 
-from src.data import RandomSplitter
+from src.data import DataSplitter, RandomSplitter
 from src.data.common import DatasetFactory
 from src.data.factories import NICOPPCfg
 from src.data.nih import NIHChestXRayDatasetCfg
@@ -15,8 +15,8 @@
 class SplitRelay:
     defaults: list[Any] = field(default_factory=lambda: [{"ds": "celeba"}, {"split": "random"}])
 
-    ds: Any  # CdtDataset
-    split: Any
+    ds: DatasetFactory
+    split: DataSplitter
     wandb: WandbConf = field(default_factory=WandbConf)
 
     options: ClassVar[dict[str, dict[str, type]]] = {
@@ -30,7 +30,6 @@ class SplitRelay:
     }
 
     def run(self, raw_config: dict[str, Any] | None = None) -> None:
-        assert isinstance(self.ds, DatasetFactory)
         assert isinstance(self.split, RandomSplitter)
 
         ds = self.ds()

src/relay/supmatch.py

@@ -57,13 +57,13 @@ class SupMatchRelay(BaseRelay):
     alg: SupportMatching = field(default_factory=SupportMatching)
     ae: SplitAeCfg = field(default_factory=SplitAeCfg)
     ae_opt: OptimizerCfg = field(default_factory=OptimizerCfg)
-    ae_arch: Any  # AeFactory
-    ds: Any  # DatasetFactory
-    disc_arch: Any  # PredictorFactory
+    ae_arch: AeFactory
+    ds: DatasetFactory
+    disc_arch: PredictorFactory
     disc: DiscOptimizerCfg = field(default_factory=DiscOptimizerCfg)
     eval: Evaluator = field(default_factory=Evaluator)
-    labeller: Any  # Labeller
-    scorer: Any  # Scorer
+    labeller: Labeller
+    scorer: Scorer
     artifact_name: str | None = None
     """Save model weights under this name."""
 
@@ -96,12 +96,6 @@ class SupMatchRelay(BaseRelay):
     }
 
     def run(self, raw_config: dict[str, Any] | None = None) -> float | None:
-        assert isinstance(self.ae_arch, AeFactory)
-        assert isinstance(self.disc_arch, PredictorFactory)
-        assert isinstance(self.ds, DatasetFactory)
-        assert isinstance(self.labeller, Labeller)
-        assert isinstance(self.scorer, Scorer)
-
         ds = self.ds()
         run = self.wandb.init(raw_config, (ds, self.labeller, self.ae_arch, self.disc_arch))
         dm = self.init_dm(ds, self.labeller, device=self.alg.device)