Add moon mkmmd example

research/flamby/fed_isic2019/fenda_mkmmd/client.py

@@ -23,7 +23,7 @@
 from fl4health.utils.losses import LossMeterType
 from fl4health.utils.metrics import BalancedAccuracy, Metric
 from fl4health.utils.random import set_all_random_seeds
-from research.flamby.fed_isic2019.fenda.fenda_model import FedIsic2019FendaModel
+from research.flamby.fed_isic2019.fenda_mkmmd.fenda_model import FedIsic2019FendaModel
 from research.flamby.flamby_data_utils import construct_fedisic_train_val_datasets
 
 

research/flamby/fed_isic2019/fenda_mkmmd/run_fold_experiment.slrm

@@ -69,8 +69,8 @@ MU=$7
 # Create the artifact directory
 mkdir "${ARTIFACT_DIR}"
 
-RUN_NAMES=( "Run1" "Run2")
-SEEDS=(2021 2022)
+RUN_NAMES=( "Run1" )
+SEEDS=(2021 )
 
 echo "Python Venv Path: ${VENV_PATH}"
 

research/flamby/fed_isic2019/fenda_mkmmd/run_hp_sweep.sh

@@ -27,12 +27,12 @@ VENV_PATH=$4
 
 # FedISIC LR Hyperparmeters from paper are not suitable for AdamW
 LR_VALUES=( 0.001 )
-MU_VALUES=( 0.0 0.1 1 10)
+MU_VALUES=( 0.1 1 10)
 
 SERVER_PORT=8100
 
 # Create sweep folder
-SWEEP_DIRECTORY="${ARTIFACT_DIR}hp_sweep_results_mkmmd"
+SWEEP_DIRECTORY="${ARTIFACT_DIR}hp_sweep_results_mkmmd_new"
 echo "Creating sweep folder at ${SWEEP_DIRECTORY}"
 mkdir ${SWEEP_DIRECTORY}
 
@@ -45,7 +45,7 @@ for LR_VALUE in "${LR_VALUES[@]}"; do
  mkdir "${EXPERIMENT_DIRECTORY}"
  SERVER_ADDRESS="0.0.0.0:${SERVER_PORT}"
  echo "Server Address: ${SERVER_ADDRESS}"
- SBATCH_COMMAND="research/flamby/fed_isic2019/fenda/run_fold_experiment.slrm \
+ SBATCH_COMMAND="research/flamby/fed_isic2019/fenda_mkmmd/run_fold_experiment.slrm \
  ${SERVER_CONFIG_PATH} \
  ${EXPERIMENT_DIRECTORY} \
  ${DATASET_DIR} \

research/flamby/fed_isic2019/fenda_mkmmd/server.py

@@ -11,7 +11,7 @@
 from fl4health.utils.config import load_config
 from fl4health.utils.functions import get_all_model_parameters
 from fl4health.utils.random import set_all_random_seeds
-from research.flamby.fed_isic2019.fenda.fenda_model import FedIsic2019FendaModel
+from research.flamby.fed_isic2019.fenda_mkmmd.fenda_model import FedIsic2019FendaModel
 from research.flamby.flamby_servers.personal_server import PersonalServer
 from research.flamby.utils import (
  evaluate_metrics_aggregation_fn,

research/flamby/fed_isic2019/moon_mkmmd/README.md

@@ -0,0 +1,24 @@
+### Running hyperparameter sweep
+
+To run the hyperparameter sweep you simply run the command
+
+```bash
+./research/flamby/fed_isic2019/moon/run_hp_sweep.sh \
+ path_to_config.yaml \
+ path_to_folder_for_artifacts/ \
+ path_to_folder_for_dataset/ \
+ path_to_desired_venv/
+```
+
+from the top level directory of the repository
+
+An example is something like
+``` bash
+./research/flamby/fed_isic2019/moon_mkmmd/run_hp_sweep.sh \
+ research/flamby/fed_isic2019/moon_mkmmd/config.yaml \
+ research/flamby/fed_isic2019/moon_mkmmd/ \
+ /Users/david/Desktop/FLambyDatasets/fedisic2019/ \
+ /h/demerson/vector_repositories/fl4health_env/
+```
+
+In order to manipulate the grid search being conducted, you need to change the parameters for `lr`, or `mu` in the `run_hp_sweep.sh` script directly. These represent the client-side learning rate and the contrastive loss weight respectively.

research/flamby/fed_isic2019/moon_mkmmd/client.py

@@ -0,0 +1,158 @@
+import argparse
+import os
+from logging import INFO
+from pathlib import Path
+from typing import Optional, Sequence, Tuple
+
+import flwr as fl
+import torch
+import torch.nn as nn
+from flamby.datasets.fed_isic2019 import BATCH_SIZE, LR, NUM_CLIENTS, BaselineLoss
+from flwr.common.logger import log
+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 fl4health.checkpointing.checkpointer import BestMetricTorchCheckpointer, TorchCheckpointer
+from fl4health.clients.moon_client import MoonClient
+from fl4health.utils.losses import LossMeterType
+from fl4health.utils.metrics import BalancedAccuracy, Metric
+from fl4health.utils.random import set_all_random_seeds
+from research.flamby.fed_isic2019.moon_mkmmd.moon_model import FedIsic2019MoonModel
+from research.flamby.flamby_data_utils import construct_fedisic_train_val_datasets
+
+
+class FedIsic2019MoonClient(MoonClient):
+ def __init__(
+ self,
+ data_path: Path,
+ metrics: Sequence[Metric],
+ device: torch.device,
+ client_number: int,
+ learning_rate: float,
+ loss_meter_type: LossMeterType = LossMeterType.AVERAGE,
+ mkmmd_loss_weights: Tuple[float, float] = (10, 10),
+ checkpointer: Optional[TorchCheckpointer] = None,
+ ) -> None:
+ super().__init__(
+ data_path=data_path,
+ metrics=metrics,
+ device=device,
+ loss_meter_type=loss_meter_type,
+ checkpointer=checkpointer,
+ mkmmd_loss_weights=mkmmd_loss_weights,
+ )
+ self.client_number = client_number
+ self.learning_rate: float = learning_rate
+
+ assert 0 <= client_number < NUM_CLIENTS
+ log(INFO, f"Client Name: {self.client_name}, Client Number: {self.client_number}")
+
+ def get_data_loaders(self, config: Config) -> Tuple[DataLoader, DataLoader]:
+ train_dataset, validation_dataset = construct_fedisic_train_val_datasets(
+ self.client_number, str(self.data_path)
+ )
+ train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
+ val_loader = DataLoader(validation_dataset, batch_size=BATCH_SIZE, shuffle=False)
+ return train_loader, val_loader
+
+ def get_model(self, config: Config) -> nn.Module:
+ model: nn.Module = FedIsic2019MoonModel().to(self.device)
+ return model
+
+ def get_optimizer(self, config: Config) -> Optimizer:
+ return torch.optim.AdamW(self.model.parameters(), lr=self.learning_rate)
+
+ def get_criterion(self, config: Config) -> _Loss:
+ return BaselineLoss()
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser(description="FL Client Main")
+ parser.add_argument(
+ "--artifact_dir",
+ action="store",
+ type=str,
+ help="Path to save client artifacts such as logs and model checkpoints",
+ required=True,
+ )
+ parser.add_argument(
+ "--dataset_dir",
+ action="store",
+ type=str,
+ help="Path to the preprocessed FedIsic2019 Dataset (ex. path/to/fedisic2019)",
+ required=True,
+ )
+ parser.add_argument(
+ "--run_name",
+ action="store",
+ help="Name of the run, model checkpoints will be saved under a subfolder with this name",
+ required=True,
+ )
+ parser.add_argument(
+ "--server_address",
+ action="store",
+ type=str,
+ help="Server Address for the clients to communicate with the server through",
+ default="0.0.0.0:8080",
+ )
+ parser.add_argument(
+ "--client_number",
+ action="store",
+ type=int,
+ help="Number of the client for dataset loading (should be 0-5 for FedIsic2019)",
+ required=True,
+ )
+ parser.add_argument(
+ "--learning_rate", action="store", type=float, help="Learning rate for local optimization", default=LR
+ )
+ parser.add_argument(
+ "--seed",
+ action="store",
+ type=int,
+ help="Seed for the random number generators across python, torch, and numpy",
+ required=False,
+ )
+ parser.add_argument(
+ "--mu",
+ action="store",
+ type=float,
+ help="Weight for the auxiliary losses",
+ required=False,
+ )
+ parser.add_argument(
+ "--gamma",
+ action="store",
+ type=float,
+ help="Weight for the auxiliary lossess",
+ required=False,
+ )
+ 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}")
+
+ # Set the random seed for reproducibility
+ set_all_random_seeds(args.seed)
+
+ 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)
+
+ client = FedIsic2019MoonClient(
+ data_path=Path(args.dataset_dir),
+ metrics=[BalancedAccuracy("FedIsic2019_balanced_accuracy")],
+ device=DEVICE,
+ client_number=args.client_number,
+ learning_rate=args.learning_rate,
+ checkpointer=checkpointer,
+ mkmmd_loss_weights=(args.mu, args.gamma),
+ )
+
+ fl.client.start_numpy_client(server_address=args.server_address, client=client)
+
+ # Shutdown the client gracefully
+ client.shutdown()

research/flamby/fed_isic2019/moon_mkmmd/config.yaml

@@ -0,0 +1,8 @@
+# Parameters that describe server
+n_server_rounds: 15 # The number of rounds to run FL
+
+# Parameters that describe clients
+n_clients: 6 # 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.

research/flamby/fed_isic2019/moon_mkmmd/moon_model.py

@@ -0,0 +1,83 @@
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from efficientnet_pytorch import EfficientNet
+from efficientnet_pytorch.utils import url_map
+from torch.utils import model_zoo
+
+from fl4health.model_bases.moon_base import MoonModel
+from research.flamby.utils import shutoff_batch_norm_tracking
+
+
+def from_pretrained(model_name: str, in_channels: int = 3, include_top: bool = False) -> EfficientNet:
+ # There is a bug in the EfficientNet implementation if you want to strip off the top layer of the network, but
+ # still load the pre-trained weights. So we do it ourselves here.
+ model = EfficientNet.from_name(model_name, include_top=include_top)
+ state_dict = model_zoo.load_url(url_map[model_name])
+ state_dict.pop("_fc.weight")
+ state_dict.pop("_fc.bias")
+ model.load_state_dict(state_dict, strict=False)
+ model._change_in_channels(in_channels)
+ return model
+
+
+class HeadClassifier(nn.Module):
+ """Moon head module"""
+
+ def __init__(self, stack_output_dimension: int):
+ super().__init__()
+ self.fc1 = nn.Linear(stack_output_dimension, 8)
+ self.dropout = nn.Dropout(0.2)
+
+ def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
+ x = self.dropout(input_tensor)
+ x = self.fc1(x)
+ return x
+
+
+class BaseEfficientNet(nn.Module):
+ """Moon feature extractor module
+ We use the EfficientNets architecture that many participants in the ISIC
+ competition have identified to work best.
+ See here the [reference paper](https://arxiv.org/abs/1905.11946)
+ Thank you to [Luke Melas-Kyriazi](https://github.com/lukemelas) for his
+ [pytorch reimplementation of EfficientNets]
+ (https://github.com/lukemelas/EfficientNet-PyTorch).
+ When loading the EfficientNet-B0 model, we strip off the FC layer to use the model as a feature extractor.
+ There is an option to freeze a subset of the layers to reduce the number of trainable parameters. However,
+ it is not used in the Moon experiments.
+ """
+
+ def __init__(self, frozen_blocks: Optional[int] = 13, turn_off_bn_tracking: bool = False):
+ super().__init__()
+ # include_top ensures that we just use feature extraction in the forward pass
+ self.base_model = from_pretrained("efficientnet-b0", include_top=False)
+ if frozen_blocks:
+ self.freeze_layers(frozen_blocks)
+ if turn_off_bn_tracking:
+ shutoff_batch_norm_tracking(self.base_model)
+
+ def freeze_layers(self, frozen_blocks: int) -> None:
+ # We freeze the bottom layers of the network. We always freeze the _conv_stem module, the _bn0 module and then
+ # we iterate throught the blocks freezing the specified number up to 15 (all of them)
+
+ # Freeze the first two layers
+ self.base_model._modules["_conv_stem"].requires_grad_(False)
+ self.base_model._modules["_bn0"].requires_grad_(False)
+ # Now we iterate through the block modules and freeze a certain number of them.
+ frozen_blocks = min(frozen_blocks, 15)
+ for block_index in range(frozen_blocks):
+ self.base_model._modules["_blocks"][block_index].requires_grad_(False)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x = self.base_model(x)
+ x = x.flatten(start_dim=1)
+ return x
+
+
+class FedIsic2019MoonModel(MoonModel):
+ def __init__(self, frozen_blocks: Optional[int] = None, turn_off_bn_tracking: bool = False) -> None:
+ base_module = BaseEfficientNet(frozen_blocks, turn_off_bn_tracking=turn_off_bn_tracking)
+ head_module = HeadClassifier(1280)
+ super().__init__(base_module, head_module)