Checkpoint Ablation for Flamby

fl4health/checkpointing/checkpointer.py

@@ -24,6 +24,7 @@ def __init__(self, best_checkpoint_dir: str, best_checkpoint_name: str) -> None:
 
     def maybe_checkpoint(self, model: nn.Module, _: Optional[float] = None) -> None:
         # Always checkpoint the latest model
+        log(INFO, "Saving latest checkpoint with LatestTorchCheckpointer")
         torch.save(model, self.best_checkpoint_path)
 
 

fl4health/clients/apfl_client.py

@@ -10,7 +10,7 @@
 from fl4health.clients.basic_client import BasicClient
 from fl4health.model_bases.apfl_base import ApflModule
 from fl4health.parameter_exchange.layer_exchanger import FixedLayerExchanger
-from fl4health.utils.losses import Losses, LossMeter, LossMeterType
+from fl4health.utils.losses import Losses, LossMeterType
 from fl4health.utils.metrics import Metric, MetricMeter, MetricMeterManager, MetricMeterType
 
 
@@ -24,11 +24,7 @@ def __init__(
         metric_meter_type: MetricMeterType = MetricMeterType.AVERAGE,
         checkpointer: Optional[TorchCheckpointer] = None,
     ) -> None:
-        super(BasicClient, self).__init__(data_path, device)
-        self.metrics = metrics
-        self.checkpointer = checkpointer
-        self.train_loss_meter = LossMeter.get_meter_by_type(loss_meter_type)
-        self.val_loss_meter = LossMeter.get_meter_by_type(loss_meter_type)
+        super().__init__(data_path, metrics, device, loss_meter_type, metric_meter_type, checkpointer)
 
         # Define mapping from prediction key to meter to pass to MetricMeterManager constructor for train and val
         train_key_to_meter_map = {

fl4health/clients/basic_client.py

@@ -152,7 +152,6 @@ def _handle_reporting(
         metric_dict: Dict[str, Scalar],
         current_round: Optional[int] = None,
     ) -> None:
-
         # If reporter is None we do not report to wandb and return
         if self.wandb_reporter is None:
             return

fl4health/clients/fenda_client.py

@@ -23,9 +23,9 @@ def __init__(
         metric_meter_type: MetricMeterType = MetricMeterType.AVERAGE,
         checkpointer: Optional[TorchCheckpointer] = None,
         temperature: Optional[float] = 0.5,
-        perfcl_loss_weights: Optional[Tuple[float, float]] = (0.0, 0.0),
-        cos_sim_loss_weight: Optional[float] = 0.0,
-        contrastive_loss_weight: Optional[float] = 0.0,
+        perfcl_loss_weights: Optional[Tuple[float, float]] = None,
+        cos_sim_loss_weight: Optional[float] = None,
+        contrastive_loss_weight: Optional[float] = None,
     ) -> None:
         super().__init__(
             data_path=data_path,
@@ -68,7 +68,6 @@ def get_parameter_exchanger(self, config: Config) -> ParameterExchanger:
         return FixedLayerExchanger(self.model.layers_to_exchange())
 
     def predict(self, input: torch.Tensor) -> Dict[str, torch.Tensor]:
-
         preds = self.model(input)
 
         if self.contrastive_loss_weight or self.perfcl_loss_weights:
@@ -88,7 +87,6 @@ def predict(self, input: torch.Tensor) -> Dict[str, torch.Tensor]:
         return preds
 
     def get_parameters(self, config: Config) -> NDArrays:
-
         # Save the parameters of the old model
         assert isinstance(self.model, FendaModel)
         if self.contrastive_loss_weight or self.perfcl_loss_weights:
@@ -98,7 +96,6 @@ def get_parameters(self, config: Config) -> NDArrays:
         return super().get_parameters(config)
 
     def set_parameters(self, parameters: NDArrays, config: Config) -> None:
-
         # Set the parameters of the model
         super().set_parameters(parameters, config)
 
@@ -181,7 +178,6 @@ def get_perfcl_loss(
         return contrastive_loss_minimize, contrastive_loss_maximize
 
     def compute_loss(self, preds: Dict[str, torch.Tensor], target: torch.Tensor) -> Losses:
-
         loss = self.criterion(preds["prediction"], target)
         total_loss = loss
         additional_losses = {}

research/flamby/fed_heart_disease/apfl/client.py

@@ -13,7 +13,11 @@
 from torch.optim import Optimizer
 from torch.utils.data import DataLoader
 
-from fl4health.checkpointing.checkpointer import BestMetricTorchCheckpointer, TorchCheckpointer
+from fl4health.checkpointing.checkpointer import (
+    BestMetricTorchCheckpointer,
+    LatestTorchCheckpointer,
+    TorchCheckpointer,
+)
 from fl4health.clients.apfl_client import ApflClient
 from fl4health.model_bases.apfl_base import ApflModule
 from fl4health.utils.losses import LossMeterType
@@ -111,17 +115,28 @@ def get_criterion(self, config: Config) -> _Loss:
     parser.add_argument(
         "--alpha_learning_rate", action="store", type=float, help="Learning rate for the APFL alpha", default=0.01
     )
+    parser.add_argument(
+        "--no_federated_checkpointing",
+        action="store_true",
+        help="boolean to indicate whether we're evaluating an APFL model or not, as those model have special args",
+    )
     args = parser.parse_args()
 
     DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     log(INFO, f"Device to be used: {DEVICE}")
     log(INFO, f"Server Address: {args.server_address}")
     log(INFO, f"Learning Rate: {args.learning_rate}")
     log(INFO, f"Alpha Learning Rate: {args.alpha_learning_rate}")
+    log(INFO, f"Performing Federated Checkpointing: {not args.no_federated_checkpointing}")
 
+    federated_checkpointing = not args.no_federated_checkpointing
     checkpoint_dir = os.path.join(args.artifact_dir, args.run_name)
     checkpoint_name = f"client_{args.client_number}_best_model.pkl"
-    checkpointer = BestMetricTorchCheckpointer(checkpoint_dir, checkpoint_name, maximize=False)
+    checkpointer = (
+        BestMetricTorchCheckpointer(checkpoint_dir, checkpoint_name, maximize=False)
+        if federated_checkpointing
+        else LatestTorchCheckpointer(checkpoint_dir, checkpoint_name)
+    )
 
     client = FedHeartDiseaseApflClient(
         data_path=args.dataset_dir,

research/flamby/fed_heart_disease/fedadam/config.yaml

@@ -6,3 +6,4 @@ n_clients: 4 # The number of clients in the FL experiment
 local_epochs: 1 # The number of epochs to complete for client (NOT USED FOR FLAMBY)
 batch_size: 64 # The batch size for client training (NOT USED FOR FLAMBY)
 local_steps: 100 # The number of local training steps to perform.
+federated_checkpointing: True # Indicates whether intermediate or latest checkpointing is used on the server side

research/flamby/fed_heart_disease/fedadam/server.py

@@ -10,7 +10,7 @@
 from flwr.server.client_manager import SimpleClientManager
 from flwr.server.strategy import FedAdam
 
-from fl4health.checkpointing.checkpointer import BestMetricTorchCheckpointer
+from fl4health.checkpointing.checkpointer import BestMetricTorchCheckpointer, LatestTorchCheckpointer
 from fl4health.utils.config import load_config
 from research.flamby.flamby_servers.full_exchange_server import FullExchangeServer
 from research.flamby.utils import (
@@ -34,7 +34,13 @@ def main(
 
     checkpoint_dir = os.path.join(checkpoint_stub, run_name)
     checkpoint_name = "server_best_model.pkl"
-    checkpointer = BestMetricTorchCheckpointer(checkpoint_dir, checkpoint_name)
+    federated_checkpointing: bool = config.get("federated_checkpointing", True)
+    log(INFO, f"Performing Federated Checkpointing: {federated_checkpointing}")
+    checkpointer = (
+        BestMetricTorchCheckpointer(checkpoint_dir, checkpoint_name)
+        if federated_checkpointing
+        else LatestTorchCheckpointer(checkpoint_dir, checkpoint_name)
+    )
 
     client_manager = SimpleClientManager()
     model = Baseline()
@@ -62,7 +68,9 @@ def main(
         config=fl.server.ServerConfig(num_rounds=config["n_server_rounds"]),
     )
 
-    log(INFO, f"Best Aggregated (Weighted) Loss seen by the Server: \n{checkpointer.best_metric}")
+    if federated_checkpointing:
+        assert isinstance(checkpointer, BestMetricTorchCheckpointer)
+        log(INFO, f"Best Aggregated (Weighted) Loss seen by the Server: \n{checkpointer.best_metric}")
 
     # Shutdown the server gracefully
     server.shutdown()

research/flamby/fed_heart_disease/fedavg/config.yaml

@@ -6,3 +6,4 @@ n_clients: 4 # The number of clients in the FL experiment
 local_epochs: 1 # The number of epochs to complete for client (NOT USED FOR FLAMBY)
 batch_size: 64 # The batch size for client training (NOT USED FOR FLAMBY)
 local_steps: 100 # The number of local training steps to perform.
+federated_checkpointing: True # Indicates whether intermediate or latest checkpointing is used on the server side

research/flamby/fed_heart_disease/fedavg/server.py

@@ -10,7 +10,7 @@
 from flwr.server.client_manager import SimpleClientManager
 from flwr.server.strategy import FedAvg
 
-from fl4health.checkpointing.checkpointer import BestMetricTorchCheckpointer
+from fl4health.checkpointing.checkpointer import BestMetricTorchCheckpointer, LatestTorchCheckpointer
 from fl4health.utils.config import load_config
 from research.flamby.flamby_servers.full_exchange_server import FullExchangeServer
 from research.flamby.utils import (
@@ -32,7 +32,13 @@ def main(config: Dict[str, Any], server_address: str, checkpoint_stub: str, run_
 
     checkpoint_dir = os.path.join(checkpoint_stub, run_name)
     checkpoint_name = "server_best_model.pkl"
-    checkpointer = BestMetricTorchCheckpointer(checkpoint_dir, checkpoint_name)
+    federated_checkpointing: bool = config.get("federated_checkpointing", True)
+    log(INFO, f"Performing Federated Checkpointing: {federated_checkpointing}")
+    checkpointer = (
+        BestMetricTorchCheckpointer(checkpoint_dir, checkpoint_name)
+        if federated_checkpointing
+        else LatestTorchCheckpointer(checkpoint_dir, checkpoint_name)
+    )
 
     client_manager = SimpleClientManager()
     model = Baseline()
@@ -60,7 +66,9 @@ def main(config: Dict[str, Any], server_address: str, checkpoint_stub: str, run_
         config=fl.server.ServerConfig(num_rounds=config["n_server_rounds"]),
     )
 
-    log(INFO, f"Best Aggregated (Weighted) Loss seen by the Server: \n{checkpointer.best_metric}")
+    if federated_checkpointing:
+        assert isinstance(checkpointer, BestMetricTorchCheckpointer)
+        log(INFO, f"Best Aggregated (Weighted) Loss seen by the Server: \n{checkpointer.best_metric}")
 
     # Shutdown the server gracefully
     server.shutdown()

research/flamby/fed_heart_disease/fedprox/config.yaml

@@ -6,3 +6,4 @@ n_clients: 4 # The number of clients in the FL experiment
 local_epochs: 1 # The number of epochs to complete for client (NOT USED FOR FLAMBY)
 batch_size: 64 # The batch size for client training (NOT USED FOR FLAMBY)
 local_steps: 100 # The number of local training steps to perform.
+federated_checkpointing: True # Indicates whether intermediate or latest checkpointing is used on the server side

research/flamby/fed_heart_disease/fedprox/server.py

@@ -9,7 +9,7 @@
 from flwr.common.logger import log
 from flwr.server.client_manager import SimpleClientManager
 
-from fl4health.checkpointing.checkpointer import BestMetricTorchCheckpointer
+from fl4health.checkpointing.checkpointer import BestMetricTorchCheckpointer, LatestTorchCheckpointer
 from fl4health.strategies.fedprox import FedProx
 from fl4health.utils.config import load_config
 from research.flamby.flamby_servers.fedprox_server import FedProxServer
@@ -32,7 +32,13 @@ def main(config: Dict[str, Any], server_address: str, mu: float, checkpoint_stub
 
     checkpoint_dir = os.path.join(checkpoint_stub, run_name)
     checkpoint_name = "server_best_model.pkl"
-    checkpointer = BestMetricTorchCheckpointer(checkpoint_dir, checkpoint_name)
+    federated_checkpointing: bool = config.get("federated_checkpointing", True)
+    log(INFO, f"Performing Federated Checkpointing: {federated_checkpointing}")
+    checkpointer = (
+        BestMetricTorchCheckpointer(checkpoint_dir, checkpoint_name)
+        if federated_checkpointing
+        else LatestTorchCheckpointer(checkpoint_dir, checkpoint_name)
+    )
 
     client_manager = SimpleClientManager()
     model = Baseline()
@@ -61,7 +67,9 @@ def main(config: Dict[str, Any], server_address: str, mu: float, checkpoint_stub
         config=fl.server.ServerConfig(num_rounds=config["n_server_rounds"]),
     )
 
-    log(INFO, f"Best Aggregated (Weighted) Loss seen by the Server: \n{checkpointer.best_metric}")
+    if federated_checkpointing:
+        assert isinstance(checkpointer, BestMetricTorchCheckpointer)
+        log(INFO, f"Best Aggregated (Weighted) Loss seen by the Server: \n{checkpointer.best_metric}")
 
     # Shutdown the server gracefully
     server.shutdown()

research/flamby/fed_heart_disease/fenda/client.py

@@ -14,7 +14,11 @@
 from torch.optim import Optimizer
 from torch.utils.data import DataLoader
 
-from fl4health.checkpointing.checkpointer import BestMetricTorchCheckpointer, TorchCheckpointer
+from fl4health.checkpointing.checkpointer import (
+    BestMetricTorchCheckpointer,
+    LatestTorchCheckpointer,
+    TorchCheckpointer,
+)
 from fl4health.clients.fenda_client import FendaClient
 from fl4health.utils.losses import LossMeterType
 from fl4health.utils.metrics import Accuracy, Metric, MetricMeterType
@@ -118,16 +122,27 @@ def get_criterion(self, config: Config) -> _Loss:
     parser.add_argument("--cos_sim_loss", action="store_true", help="Activate Cosine Similarity loss")
     parser.add_argument("--contrastive_loss", action="store_true", help="Activate Contrastive loss")
     parser.add_argument("--perfcl_loss", action="store_true", help="Activate PerFCL loss")
+    parser.add_argument(
+        "--no_federated_checkpointing",
+        action="store_true",
+        help="boolean to indicate whether we're evaluating an APFL model or not, as those model have special args",
+    )
     args = parser.parse_args()
 
     DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     log(INFO, f"Device to be used: {DEVICE}")
     log(INFO, f"Server Address: {args.server_address}")
     log(INFO, f"Learning Rate: {args.learning_rate}")
+    log(INFO, f"Performing Federated Checkpointing: {not args.no_federated_checkpointing}")
 
+    federated_checkpointing = not args.no_federated_checkpointing
     checkpoint_dir = os.path.join(args.artifact_dir, args.run_name)
     checkpoint_name = f"client_{args.client_number}_best_model.pkl"
-    checkpointer = BestMetricTorchCheckpointer(checkpoint_dir, checkpoint_name, maximize=False)
+    checkpointer = (
+        BestMetricTorchCheckpointer(checkpoint_dir, checkpoint_name, maximize=False)
+        if federated_checkpointing
+        else LatestTorchCheckpointer(checkpoint_dir, checkpoint_name)
+    )
 
     client = FedHeartDiseaseFendaClient(
         data_path=Path(args.dataset_dir),

research/flamby/fed_heart_disease/scaffold/config.yaml

@@ -6,3 +6,4 @@ n_clients: 4 # The number of clients in the FL experiment
 local_epochs: 1 # The number of epochs to complete for client (NOT USED FOR FLAMBY)
 batch_size: 64 # The batch size for client training (NOT USED FOR FLAMBY)
 local_steps: 100 # The number of local training steps to perform.
+federated_checkpointing: True # Indicates whether intermediate or latest checkpointing is used on the server side

research/flamby/fed_heart_disease/scaffold/server.py

@@ -8,7 +8,7 @@
 from flamby.datasets.fed_heart_disease import Baseline
 from flwr.common.logger import log
 
-from fl4health.checkpointing.checkpointer import BestMetricTorchCheckpointer
+from fl4health.checkpointing.checkpointer import BestMetricTorchCheckpointer, LatestTorchCheckpointer
 from fl4health.client_managers.fixed_without_replacement_manager import FixedSamplingByFractionClientManager
 from fl4health.strategies.scaffold import Scaffold
 from fl4health.utils.config import load_config
@@ -34,7 +34,13 @@ def main(
 
     checkpoint_dir = os.path.join(checkpoint_stub, run_name)
     checkpoint_name = "server_best_model.pkl"
-    checkpointer = BestMetricTorchCheckpointer(checkpoint_dir, checkpoint_name)
+    federated_checkpointing: bool = config.get("federated_checkpointing", True)
+    log(INFO, f"Performing Federated Checkpointing: {federated_checkpointing}")
+    checkpointer = (
+        BestMetricTorchCheckpointer(checkpoint_dir, checkpoint_name)
+        if federated_checkpointing
+        else LatestTorchCheckpointer(checkpoint_dir, checkpoint_name)
+    )
 
     client_manager = FixedSamplingByFractionClientManager()
     model = Baseline()
@@ -65,7 +71,9 @@ def main(
         config=fl.server.ServerConfig(num_rounds=config["n_server_rounds"]),
     )
 
-    log(INFO, f"Best Aggregated (Weighted) Loss seen by the Server: \n{checkpointer.best_metric}")
+    if federated_checkpointing:
+        assert isinstance(checkpointer, BestMetricTorchCheckpointer)
+        log(INFO, f"Best Aggregated (Weighted) Loss seen by the Server: \n{checkpointer.best_metric}")
 
     # Shutdown the server gracefully
     server.shutdown()

research/flamby/fed_ixi/apfl/client.py

@@ -14,7 +14,11 @@
 from torch.optim import Optimizer
 from torch.utils.data import DataLoader
 
-from fl4health.checkpointing.checkpointer import BestMetricTorchCheckpointer, TorchCheckpointer
+from fl4health.checkpointing.checkpointer import (
+    BestMetricTorchCheckpointer,
+    LatestTorchCheckpointer,
+    TorchCheckpointer,
+)
 from fl4health.clients.apfl_client import ApflClient
 from fl4health.model_bases.apfl_base import ApflModule
 from fl4health.utils.losses import LossMeterType
@@ -50,7 +54,7 @@ def __init__(
         self.alpha_learning_rate = alpha_learning_rate
         self.client_number = client_number
 
-    def get_dataloader(self, config: Config) -> Tuple[DataLoader, DataLoader]:
+    def get_data_loaders(self, config: Config) -> Tuple[DataLoader, DataLoader]:
         train_dataset, validation_dataset = construct_fed_ixi_train_val_datasets(
             self.client_number, str(self.data_path)
         )
@@ -113,17 +117,28 @@ def get_criterion(self, config: Config) -> _Loss:
     parser.add_argument(
         "--alpha_learning_rate", action="store", type=float, help="Learning rate for the APFL alpha", default=0.01
     )
+    parser.add_argument(
+        "--no_federated_checkpointing",
+        action="store_true",
+        help="boolean to indicate whether we're evaluating an APFL model or not, as those model have special args",
+    )
     args = parser.parse_args()
 
     DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     log(INFO, f"Device to be used: {DEVICE}")
     log(INFO, f"Server Address: {args.server_address}")
     log(INFO, f"Learning Rate: {args.learning_rate}")
     log(INFO, f"Alpha Learning Rate: {args.alpha_learning_rate}")
+    log(INFO, f"Performing Federated Checkpointing: {not args.no_federated_checkpointing}")
 
+    federated_checkpointing = not args.no_federated_checkpointing
     checkpoint_dir = os.path.join(args.artifact_dir, args.run_name)
     checkpoint_name = f"client_{args.client_number}_best_model.pkl"
-    checkpointer = BestMetricTorchCheckpointer(checkpoint_dir, checkpoint_name, maximize=False)
+    checkpointer = (
+        BestMetricTorchCheckpointer(checkpoint_dir, checkpoint_name, maximize=False)
+        if federated_checkpointing
+        else LatestTorchCheckpointer(checkpoint_dir, checkpoint_name)
+    )
 
     client = FedIxiApflClient(
         data_path=Path(args.dataset_dir),

research/flamby/fed_ixi/fedadam/config.yaml

@@ -6,3 +6,4 @@ n_clients: 3 # The number of clients in the FL experiment
 local_epochs: 1 # The number of epochs to complete for client (NOT USED FOR FLAMBY)
 batch_size: 64 # The batch size for client training (NOT USED FOR FLAMBY)
 local_steps: 100 # The number of local training steps to perform.
+federated_checkpointing: True # Indicates whether intermediate or latest checkpointing is used on the server side