🎨 Rename attribute model to core_model in all wrappers

Author: alafage

tests/models/wrappers/test_batch_ensemble.py

@@ -45,8 +45,8 @@ def test_convert_layers(self) -> None:
         model = _DummyModel(in_features, out_features)
         wrapped_model = batch_ensemble(model, num_estimators, convert_layers=True)
         assert wrapped_model.num_estimators == num_estimators
-        assert isinstance(wrapped_model.model.conv, BatchConv2d)
-        assert isinstance(wrapped_model.model.fc, BatchLinear)
+        assert isinstance(wrapped_model.core_model.conv, BatchConv2d)
+        assert isinstance(wrapped_model.core_model.fc, BatchLinear)
 
     def test_forward_pass(self, img_input) -> None:
         batch_size = img_input.size(0)

tests/models/wrappers/test_mc_dropout.py

@@ -135,7 +135,7 @@ def test_mc_dropout_errors(self) -> None:
             mc_dropout(model, num_estimators=5, task="regression")
 
         with pytest.raises(ValueError, match="`num_estimators` must be strictly positive"):
-            mc_dropout(model=model, num_estimators=-1, last_layer=True, on_batch=True)
+            mc_dropout(core_model=model, num_estimators=-1, last_layer=True, on_batch=True)
 
         dropout_model = mc_dropout(model, 5)
         with pytest.raises(TypeError, match="Training mode is expected to be boolean"):

torch_uncertainty/models/classification/lenet.py

@@ -139,7 +139,7 @@ def batchensemble_lenet(
         dropout_rate=dropout_rate,
     )
     return BatchEnsemble(
-        model=model,
+        core_model=model,
         num_estimators=num_estimators,
         repeat_training_inputs=repeat_training_inputs,
         convert_layers=True,

torch_uncertainty/models/mlp.py

@@ -173,7 +173,7 @@ def _mlp(
     )
     if stochastic:
         return StochasticModel(
-            model=model, num_samples=num_samples, probabilistic=dist_family is not None
+            core_model=model, num_samples=num_samples, probabilistic=dist_family is not None
         )
     return model
 

torch_uncertainty/models/wrappers/batch_ensemble.py

@@ -8,7 +8,7 @@
 class BatchEnsemble(nn.Module):
     def __init__(
         self,
-        model: nn.Module,
+        core_model: nn.Module,
         num_estimators: int,
         repeat_training_inputs: bool = False,
         convert_layers: bool = False,
@@ -23,7 +23,7 @@ def __init__(
         ensuring that each estimator receives the correct data format.
 
         Args:
-            model (nn.Module): The BatchEnsemble model.
+            core_model (nn.Module): The BatchEnsemble model.
             num_estimators (int): Number of ensemble members.
             repeat_training_inputs (optional, bool): Whether to repeat the input batch during training.
                 If ``True``, the input batch is repeated during both training and evaluation. If ``False``,
@@ -33,37 +33,37 @@ def __init__(
                 BatchEnsemble counterparts. Default is ``False``.
 
         Raises:
-            ValueError: If neither ``BatchLinear`` nor ``BatchConv2d`` layers are found in the model at the
+            ValueError: If neither ``BatchLinear`` nor ``BatchConv2d`` layers are found in the core_model at the
                 end of initialization.
             ValueError: If ``num_estimators`` is less than or equal to ``0``.
             ValueError: If ``convert_layers=True`` and neither ``nn.Linear`` nor ``nn.Conv2d`` layers are
-                found in the model.
+                found in the core_model.
 
         Warning:
             If ``convert_layers==True``, the wrapper will attempt to convert all ``nn.Linear`` and ``nn.Conv2d``
-            layers in the model to their BatchEnsemble counterparts. If the model contains other types of
+            layers in the core_model to their BatchEnsemble counterparts. If the core_model contains other types of
             layers, the conversion won't happen for these layers. If don't have any ``nn.Linear`` or ``nn.Conv2d``
-            layers in the model, the wrapper will raise an error during conversion.
+            layers in the core_model, the wrapper will raise an error during conversion.
 
         Warning:
             If ``repeat_training_inputs==True`` and you want to use one of the ``torch_uncertainty.routines``
             for training, be sure to set ``format_batch_fn=RepeatTarget(num_repeats=num_estimators)`` when
             initializing the routine.
 
         Example:
-            >>> model = nn.Sequential(nn.Linear(10, 5), nn.ReLU(), nn.Linear(5, 2))
-            >>> model = BatchEnsemble(model, num_estimators=4, convert_layers=True)
+            >>> core_model = nn.Sequential(nn.Linear(10, 5), nn.ReLU(), nn.Linear(5, 2))
+            >>> model = BatchEnsemble(core_model, num_estimators=4, convert_layers=True)
             >>> model
             BatchEnsemble(
-            (model): Sequential(
+            (core_model): Sequential(
                 (0): BatchLinear(in_features=10, out_features=5, num_estimators=4)
                 (1): ReLU()
                 (2): BatchLinear(in_features=5, out_features=2, num_estimators=4)
             )
             )
         """
         super().__init__()
-        self.model = model
+        self.core_model = core_model
         self.num_estimators = num_estimators
         self.repeat_training_inputs = repeat_training_inputs
 
@@ -72,7 +72,7 @@ def __init__(
 
         filtered_modules = [
             module
-            for module in self.model.modules()
+            for module in self.core_model.modules()
             if isinstance(module, BatchLinear | BatchConv2d)
         ]
         _batch_ensemble_checks(filtered_modules, num_estimators)
@@ -81,22 +81,22 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Repeat the input if ``self.training==False`` or ``repeat_training_inputs==True`` and pass it through the model."""
         if not self.training or self.repeat_training_inputs:
             x = repeat(x, "b ... -> (m b) ...", m=self.num_estimators)
-        return self.model(x)
+        return self.core_model(x)
 
     def _convert_layers(self) -> None:
         """Convert the model's layers to BatchEnsemble layers."""
         no_valid_layers = True
-        for name, layer in self.model.named_modules():
+        for name, layer in self.core_model.named_modules():
             if isinstance(layer, nn.Linear):
                 setattr(
-                    self.model,
+                    self.core_model,
                     name,
                     BatchLinear.from_linear(layer, num_estimators=self.num_estimators),
                 )
                 no_valid_layers = False
             elif isinstance(layer, nn.Conv2d):
                 setattr(
-                    self.model,
+                    self.core_model,
                     name,
                     BatchConv2d.from_conv2d(layer, num_estimators=self.num_estimators),
                 )
@@ -121,15 +121,15 @@ def _batch_ensemble_checks(filtered_modules: list[nn.Module], num_estimators: in
 
 
 def batch_ensemble(
-    model: nn.Module,
+    core_model: nn.Module,
     num_estimators: int,
     repeat_training_inputs: bool = False,
     convert_layers: bool = False,
 ) -> BatchEnsemble:
     """BatchEnsemble wrapper for a model.
 
     Args:
-        model (nn.Module): model to wrap
+        core_model (nn.Module): model to wrap
         num_estimators (int): number of ensemble members
         repeat_training_inputs (bool, optional): whether to repeat the input batch during training.
             If ``True``, the input batch is repeated during both training and evaluation. If ``False``,
@@ -139,10 +139,10 @@ def batch_ensemble(
             BatchEnsemble counterparts. Default is ``False``.
 
     Returns:
-        BatchEnsemble: BatchEnsemble wrapper for the model
+        BatchEnsemble: BatchEnsemble wrapper for the :attr:`core_model`
     """
     return BatchEnsemble(
-        model=model,
+        core_model=core_model,
         num_estimators=num_estimators,
         repeat_training_inputs=repeat_training_inputs,
         convert_layers=convert_layers,

torch_uncertainty/models/wrappers/checkpoint_collector.py

@@ -7,7 +7,7 @@
 class CheckpointCollector(nn.Module):
     def __init__(
         self,
-        model: nn.Module,
+        core_model: nn.Module,
         cycle_start: int | None = None,
         cycle_length: int | None = None,
         save_schedule: list[int] | None = None,
@@ -21,7 +21,7 @@ def __init__(
         as implemented in TorchUncertainty.
 
         Args:
-            model (nn.Module): The model to train and ensemble.
+            core_model (nn.Module): The model to train and ensemble.
             cycle_start (int): Epoch to start ensembling. Defaults to ``None``.
             cycle_length (int): Number of epochs between model collections. Defaults to ``None``.
             save_schedule (list[int] | None): The epochs at which to save the model. Defaults to ``None``.
@@ -52,7 +52,7 @@ def __init__(
                 f"The combination of arguments: cycle_start: {cycle_start}, cycle_length: {cycle_length}, save_schedule: {save_schedule} is not known."
             )
 
-        self.core_model = model
+        self.core_model = core_model
         self.cycle_start = cycle_start
         self.cycle_length = cycle_length
         self.save_schedule = save_schedule

torch_uncertainty/models/wrappers/deep_ensembles.py

@@ -9,13 +9,13 @@
 class _DeepEnsembles(nn.Module):
     def __init__(
         self,
-        models: list[nn.Module],
+        core_models: list[nn.Module],
         store_on_cpu: bool = False,
     ) -> None:
         """Create a classification deep ensembles from a list of models."""
         super().__init__()
-        self.core_models = nn.ModuleList(models)
-        self.num_estimators = len(models)
+        self.core_models = nn.ModuleList(core_models)
+        self.num_estimators = len(core_models)
         self.store_on_cpu = store_on_cpu
 
     def forward(self, x: Tensor) -> Tensor:
@@ -52,11 +52,11 @@ class _RegDeepEnsembles(_DeepEnsembles):
     def __init__(
         self,
         probabilistic: bool,
-        models: list[nn.Module],
+        core_models: list[nn.Module],
         store_on_cpu: bool = False,
     ) -> None:
         """Create a regression deep ensembles from a list of models."""
-        super().__init__(models=models, store_on_cpu=store_on_cpu)
+        super().__init__(core_models=core_models, store_on_cpu=store_on_cpu)
         self.probabilistic = probabilistic
 
     def forward(self, x: Tensor) -> Tensor | dict[str, Tensor]:
@@ -87,7 +87,7 @@ def forward(self, x: Tensor) -> Tensor | dict[str, Tensor]:
 
 
 def deep_ensembles(
-    models: list[nn.Module] | nn.Module,
+    core_models: list[nn.Module] | nn.Module,
     num_estimators: int | None = None,
     task: Literal[
         "classification", "regression", "segmentation", "pixel_regression"
@@ -101,12 +101,12 @@ def deep_ensembles(
     """Build a Deep Ensembles out of the original models.
 
     Args:
-        models (list[nn.Module] | nn.Module): The model to be ensembled.
+        core_models (list[nn.Module] | nn.Module): The model to be ensembled.
         num_estimators (int | None): The number of estimators in the ensemble.
         task (Literal[``"classification"``, ``"regression"``, ``"segmentation"``, ``"pixel_regression"``]): The model task. Defaults to ``"classification"``.
         probabilistic (bool): Whether the regression model is probabilistic.
         reset_model_parameters (bool): Whether to reset the model parameters
-            when :attr:models is a module or a list of length 1. Defaults to ``True``.
+            when :attr:core_models is a module or a list of length 1. Defaults to ``True``.
         store_on_cpu (bool): Whether to store the models on CPU. Defaults to ``False``.
             This is useful for large models that do not fit in GPU memory. Only one
             model will be stored on GPU at a time during forward. The rest will be stored on CPU.
@@ -140,26 +140,28 @@ def deep_ensembles(
             <https://arxiv.org/abs/1612.01474>`_.
 
     """
-    if isinstance(models, list) and len(models) == 0:
+    if isinstance(core_models, list) and len(core_models) == 0:
         raise ValueError("Models must not be an empty list.")
-    if (isinstance(models, list) and len(models) == 1) or isinstance(models, nn.Module):
+    if (isinstance(core_models, list) and len(core_models) == 1) or isinstance(
+        core_models, nn.Module
+    ):
         if num_estimators is None:
             raise ValueError("if models is a module, num_estimators must be specified.")
         if num_estimators < 2:
             raise ValueError(f"num_estimators must be at least 2. Got {num_estimators}.")
 
-        if isinstance(models, list):
-            models = models[0]
+        if isinstance(core_models, list):
+            core_models = core_models[0]
 
-        models = [copy.deepcopy(models) for _ in range(num_estimators)]
+        core_models = [copy.deepcopy(core_models) for _ in range(num_estimators)]
 
         if reset_model_parameters:
-            for model in models:
+            for model in core_models:
                 for layer in model.modules():
                     if hasattr(layer, "reset_parameters"):
                         layer.reset_parameters()
 
-    elif isinstance(models, list) and len(models) > 1 and num_estimators is not None:
+    elif isinstance(core_models, list) and len(core_models) > 1 and num_estimators is not None:
         raise ValueError("num_estimators must be None if you provided a non-singleton list.")
 
     if ckpt_paths is not None:  # coverage: ignore
@@ -175,11 +177,11 @@ def deep_ensembles(
             if len(ckpt_paths) == 0:
                 raise ValueError("No checkpoint files found in the directory.")
 
-        if len(models) != len(ckpt_paths):
+        if len(core_models) != len(ckpt_paths):
             raise ValueError(
                 "The number of models and the number of checkpoint paths must be the same."
             )
-        for model, ckpt_path in zip(models, ckpt_paths, strict=True):
+        for model, ckpt_path in zip(core_models, ckpt_paths, strict=True):
             if isinstance(ckpt_path, str | Path):
                 loaded_data = torch.load(ckpt_path, map_location="cpu")
                 if "state_dict" in loaded_data:
@@ -198,12 +200,12 @@ def deep_ensembles(
 
     match task:
         case "classification" | "segmentation":
-            return _DeepEnsembles(models=models, store_on_cpu=store_on_cpu)
+            return _DeepEnsembles(core_models=core_models, store_on_cpu=store_on_cpu)
         case "regression" | "pixel_regression":
             if probabilistic is None:
                 raise ValueError("probabilistic must be specified for regression models.")
             return _RegDeepEnsembles(
-                probabilistic=probabilistic, models=models, store_on_cpu=store_on_cpu
+                probabilistic=probabilistic, core_models=core_models, store_on_cpu=store_on_cpu
             )
         case _:
             raise ValueError(f"Unknown task: {task}.")

torch_uncertainty/models/wrappers/ema.py

@@ -6,7 +6,7 @@
 class EMA(nn.Module):
     def __init__(
         self,
-        model: nn.Module,
+        core_model: nn.Module,
         momentum: float,
     ) -> None:
         """Exponential Moving Average (EMA).
@@ -18,13 +18,13 @@ def __init__(
         The EMA model is regularly updated with the inner-model and used at evaluation time.
 
         Args:
-            model (nn.Module): The model to train and ensemble.
+            core_model (nn.Module): The model to train and ensemble.
             momentum (float): The momentum of the moving average.
         """
         super().__init__()
         _ema_checks(momentum)
-        self.core_model = model
-        self.ema_model = copy.deepcopy(model)
+        self.core_model = core_model
+        self.ema_model = copy.deepcopy(core_model)
         self.momentum = momentum
         self.remainder = 1 - momentum