Creating FedBN client and adding test

examples/fedbn_example/client.py

@@ -13,7 +13,7 @@
 from torch.utils.data import DataLoader
 
 from examples.models.cnn_model import MnistNetWithBnAndFrozen, SkinCancerNet
-from fl4health.clients.basic_client import BasicClient
+from fl4health.clients.fedbn_client import FedBnClient
 from fl4health.datasets.skin_cancer.load_data import load_skin_cancer_data
 from fl4health.metrics import Accuracy
 from fl4health.metrics.base_metrics import Metric
@@ -24,7 +24,7 @@
 from fl4health.utils.sampler import DirichletLabelBasedSampler
 
 
-class MnistFedBNClient(BasicClient):
+class MnistFedBnClient(FedBnClient):
     def get_data_loaders(self, config: Config) -> tuple[DataLoader, DataLoader]:
         batch_size = narrow_dict_type(config, "batch_size", int)
         sampler = DirichletLabelBasedSampler(list(range(10)), sample_percentage=0.75, beta=1)
@@ -45,7 +45,7 @@ def get_parameter_exchanger(self, config: Config) -> ParameterExchanger:
         return LayerExchangerWithExclusions(self.model, {nn.BatchNorm2d})
 
 
-class SkinCancerFedBNClient(BasicClient):
+class SkinCancerFedBNClient(FedBnClient):
     def __init__(self, data_path: Path, metrics: Sequence[Metric], device: torch.device, dataset_name: str):
         super().__init__(data_path, metrics, device)
         self.dataset_name = dataset_name
@@ -95,9 +95,9 @@ def get_parameter_exchanger(self, config: Config) -> ParameterExchanger:
     log(INFO, f"Server Address: {args.server_address}")
 
     if args.dataset_name in ["Barcelona", "Rosendahl", "Vienna", "UFES", "Canada"]:
-        client: BasicClient = SkinCancerFedBNClient(data_path, [Accuracy()], device, args.dataset_name)
+        client: FedBnClient = SkinCancerFedBNClient(data_path, [Accuracy()], device, args.dataset_name)
     elif args.dataset_name == "mnist":
-        client = MnistFedBNClient(data_path, [Accuracy()], device)
+        client = MnistFedBnClient(data_path, [Accuracy()], device)
     else:
         raise ValueError(
             "Unsupported dataset name. Please choose from 'Barcelona', 'Rosendahl', \

fl4health/clients/fedbn_client.py

@@ -0,0 +1,27 @@
+from flwr.common.typing import Config
+
+from fl4health.clients.basic_client import BasicClient
+from fl4health.parameter_exchange.layer_exchanger import LayerExchangerWithExclusions
+
+
+class FedBnClient(BasicClient):
+    """
+    This class serves as a sparse interface for clients aiming to leverage the FedBN method
+    (https://arxiv.org/abs/2102.07623) or any other approach that excludes specific types of model layers during
+    parameter exchange. This class simply ensures that the user has overridden the `get_parameter_exchanger` properly.
+
+    For example, in FedBN, batch normalization layers are excluded from exchange with the server
+    but all other layers flow through and are aggregated via whatever strategy the server is implementing. An example
+    of this where one wants to exclude 2D batch normalization layers during exchange is
+    `LayerExchangerWithExclusions(self.model, {nn.BatchNorm2d})`, where the model is provided so that the exchanger
+    can identify the appropriate layers to leave out.
+    """
+
+    def setup_client(self, config: Config) -> None:
+        super().setup_client(config=config)
+        assert isinstance(self.parameter_exchanger, LayerExchangerWithExclusions), (
+            "For FedBnClients the parameter exchanger must be of type LayerExchangerWithExclusions "
+            f"but got {type(self.parameter_exchanger)}. If you haven't already, override the get_parameter_exchanger "
+            "function in your class."
+        )
+        return super().setup_client(config)

fl4health/strategies/fedpm.py

@@ -110,7 +110,7 @@ def aggregate_bayesian(self, results: list[tuple[NDArrays, int]]) -> dict[str, N
 
         In this case, the updates performed are:
 
-        .. code:: python
+        .. code-block:: python
 
             alpha_new = alpha + M
 

pyproject.toml

@@ -33,6 +33,8 @@ wandb = "^0.18.0"
 acvl_utils = "0.2" # Pin as it was causing an issue with nnunet (ModuleNotFoundError: No module named 'blosc2')
 scikit-learn = "1.5.0" # Pin as it was causing issues with nnunet
 peft = "^0.14.0"
+jupyter-core = "^5.8.1"
+fastapi = "^0.115.12"
 
 # Problematic grpcio versions cause issues, should be fixed in next flwr update
 # See https://github.com/adap/flower/pull/3853

tests/clients/test_fedbn_client.py

@@ -0,0 +1,66 @@
+from pathlib import Path
+
+import pytest
+import torch
+import torch.nn as nn
+from flwr.common.typing import Config
+from torch.nn.modules.loss import _Loss
+from torch.optim import Optimizer
+from torch.utils.data import DataLoader
+from torch.utils.data.dataset import TensorDataset
+
+from fl4health.clients.fedbn_client import FedBnClient
+from fl4health.metrics import Accuracy
+from fl4health.parameter_exchange.full_exchanger import FullParameterExchanger
+from fl4health.parameter_exchange.layer_exchanger import LayerExchangerWithExclusions
+from fl4health.parameter_exchange.parameter_exchanger_base import ParameterExchanger
+from tests.test_utils.models_for_test import ToyConvNet
+
+
+class GoodClientForTest(FedBnClient):
+    def get_parameter_exchanger(self, config: Config) -> ParameterExchanger:
+        return LayerExchangerWithExclusions(ToyConvNet(include_bn=True), {nn.BatchNorm1d})
+
+    def get_data_loaders(self, config: Config) -> tuple[DataLoader, DataLoader]:
+        train_loader = DataLoader(TensorDataset(torch.ones((4, 4)), torch.ones((4))))
+        val_loader = DataLoader(TensorDataset(torch.ones((4, 4)), torch.ones((4))))
+        return train_loader, val_loader
+
+    def get_criterion(self, config: Config) -> _Loss:
+        return torch.nn.CrossEntropyLoss()
+
+    def get_optimizer(self, config: Config) -> Optimizer:
+        return torch.optim.SGD(self.model.parameters(), lr=0.001, momentum=0.9)
+
+    def get_model(self, config: Config) -> nn.Module:
+        return ToyConvNet().to(self.device)
+
+
+class BadClientForTest(FedBnClient):
+    def get_parameter_exchanger(self, config: Config) -> ParameterExchanger:
+        return FullParameterExchanger()
+
+    def get_data_loaders(self, config: Config) -> tuple[DataLoader, DataLoader]:
+        train_loader = DataLoader(TensorDataset(torch.ones((4, 4)), torch.ones((4))))
+        val_loader = DataLoader(TensorDataset(torch.ones((4, 4)), torch.ones((4))))
+        return train_loader, val_loader
+
+    def get_criterion(self, config: Config) -> _Loss:
+        return torch.nn.CrossEntropyLoss()
+
+    def get_optimizer(self, config: Config) -> Optimizer:
+        return torch.optim.SGD(self.model.parameters(), lr=0.001, momentum=0.9)
+
+    def get_model(self, config: Config) -> nn.Module:
+        return ToyConvNet().to(self.device)
+
+
+def test_instance_level_client_with_changes() -> None:
+    # Creating client in the right way shouldn't throw an error
+    good_client = GoodClientForTest(data_path=Path(""), metrics=[Accuracy()], device=torch.device("cpu"))
+    good_client.setup_client({})
+    # We should throw an assertion error here because we're trying to use a FedBnClient with the wrong kind of
+    # parameter exchanger.
+    with pytest.raises(AssertionError):
+        bad_client = BadClientForTest(data_path=Path(""), metrics=[Accuracy()], device=torch.device("cpu"))
+        bad_client.setup_client({})

tests/mixins/personalized/test_ditto.py

@@ -24,9 +24,7 @@
     make_it_personal,
 )
 from fl4health.parameter_exchange.packing_exchanger import FullParameterExchangerWithPacking
-from fl4health.parameter_exchange.parameter_packer import (
-    ParameterPackerAdaptiveConstraint,
-)
+from fl4health.parameter_exchange.parameter_packer import ParameterPackerAdaptiveConstraint
 from fl4health.utils.losses import TrainingLosses
 from fl4health.utils.typing import TorchFeatureType, TorchInputType, TorchPredType
 

tests/mixins/personalized/test_utils.py

@@ -13,9 +13,7 @@
 from fl4health.metrics import Accuracy
 from fl4health.mixins.personalized.utils import ensure_protocol_compliance
 from fl4health.parameter_exchange.packing_exchanger import FullParameterExchangerWithPacking
-from fl4health.parameter_exchange.parameter_packer import (
-    ParameterPackerAdaptiveConstraint,
-)
+from fl4health.parameter_exchange.parameter_packer import ParameterPackerAdaptiveConstraint
 
 
 def test_ensure_protocol_compliance_does_not_raise() -> None:

tests/mixins/test_adaptive_drift_constrained.py

@@ -22,9 +22,7 @@
 )
 from fl4health.mixins.core_protocols import BasicClientProtocol
 from fl4health.parameter_exchange.packing_exchanger import FullParameterExchangerWithPacking
-from fl4health.parameter_exchange.parameter_packer import (
-    ParameterPackerAdaptiveConstraint,
-)
+from fl4health.parameter_exchange.parameter_packer import ParameterPackerAdaptiveConstraint
 
 
 class _TestBasicClient(BasicClient):