[Feature] Add DittoPersonalizedMixin

fl4health/mixins/core_protocols.py

@@ -7,7 +7,8 @@
 from torch.optim import Optimizer
 from torch.utils.data import DataLoader
 
-from fl4health.utils.typing import TorchFeatureType, TorchPredType, TorchTargetType
+from fl4health.utils.losses import EvaluationLosses, TrainingLosses
+from fl4health.utils.typing import TorchFeatureType, TorchInputType, TorchPredType, TorchTargetType
 
 
 @runtime_checkable
@@ -68,3 +69,35 @@ class BasicClientProtocol(BasicClientProtocolPreSetup, Protocol):
     train_loader: DataLoader
     val_loader: DataLoader
     test_loader: DataLoader | None
+    criterion: _Loss
+
+    def initialize_all_model_weights(self, parameters: NDArrays, config: Config) -> None:
+        pass  # pragma: no cover
+
+    def update_before_train(self, current_server_round: int) -> None:
+        pass  # pragma: no cover
+
+    def predict(self, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureType]:
+        pass  # pragma: no cover
+
+    def transform_target(self, target: TorchTargetType) -> TorchTargetType:
+        pass  # pragma: no cover
+
+    def compute_training_loss(
+        self,
+        preds: TorchPredType,
+        features: TorchFeatureType,
+        target: TorchTargetType,
+    ) -> TrainingLosses:
+        pass  # pragma: no cover
+
+    def validate(self, include_losses_in_metrics: bool = False) -> tuple[float, dict[str, Scalar]]:
+        pass  # pragma: no cover
+
+    def compute_evaluation_loss(
+        self,
+        preds: TorchPredType,
+        features: TorchFeatureType,
+        target: TorchTargetType,
+    ) -> EvaluationLosses:
+        pass  # pragma: no cover

fl4health/mixins/personalized/__init__.py

@@ -0,0 +1,39 @@
+from enum import Enum
+
+from fl4health.clients.basic_client import BasicClient
+from fl4health.mixins.personalized.ditto import DittoPersonalizedMixin, DittoPersonalizedProtocol
+
+
+class PersonalizedMode(Enum):
+    DITTO = "ditto"
+
+
+PersonalizedMixinRegistry = {PersonalizedMode.DITTO: DittoPersonalizedMixin}
+
+
+def make_it_personal(client_base_type: type[BasicClient], mode: PersonalizedMode) -> type[BasicClient]:
+    """A mixed class factory for converting basic clients to personalized versions."""
+    if mode == PersonalizedMode.DITTO:
+
+        return type(
+            f"Ditto{client_base_type.__name__}",
+            (
+                PersonalizedMixinRegistry[mode],
+                client_base_type,
+            ),
+            {
+                # Special flag to bypass validation
+                "_dynamically_created": True
+            },
+        )
+    else:
+        raise ValueError("Unrecognized personalized mode.")
+
+
+__all__ = [
+    "DittoPersonalizedMixin",
+    "DittoPersonalizedProtocol",
+    "PersonalizedMode",
+    "PersonalizedMixinRegistry",
+    "make_it_personal",
+]