Merge pull request #80 from VectorInstitute/multi-label

bug fix in basic client

Author: yuchongzhang

examples/feature_alignment_example/client.py

@@ -1,7 +1,7 @@
 import argparse
 from logging import INFO
 from pathlib import Path
-from typing import Sequence, Tuple
+from typing import List, Sequence, Tuple, Union
 
 import flwr as fl
 import numpy as np
@@ -22,9 +22,14 @@
 
 class Mimic3TabularDataClient(TabularDataClient):
     def __init__(
-        self, data_path: Path, metrics: Sequence[Metric], device: torch.device, id_column: str, target_column: str
+        self,
+        data_path: Path,
+        metrics: Sequence[Metric],
+        device: torch.device,
+        id_column: str,
+        targets: Union[str, List[str]],
     ) -> None:
-        super().__init__(data_path, metrics, device, id_column, target_column)
+        super().__init__(data_path, metrics, device, id_column, targets)
 
     def get_data_loaders(self, config: Config) -> Tuple[DataLoader, DataLoader]:
         batch_size = self.narrow_config_type(config, "batch_size", int)
@@ -92,7 +97,7 @@ def get_data_frame(self, config: Config) -> pd.DataFrame:
     log(INFO, f"Server Address: {args.server_address}")
 
     # ham_id is the id column and LOSgroupNum is the target column.
-    client = Mimic3TabularDataClient(data_path, [Accuracy("accuracy")], DEVICE, "hadm_id", "LOSgroupNum")
+    client = Mimic3TabularDataClient(data_path, [Accuracy("accuracy")], DEVICE, "hadm_id", ["LOSgroupNum"])
     # This call demonstrates how the user may specify a particular sklearn pipeline for a specific feature.
     client.preset_specific_pipeline("NumNotes", MaxAbsScaler())
     fl.client.start_numpy_client(server_address=args.server_address, client=client)

fl4health/clients/basic_client.py

@@ -557,7 +557,7 @@ def setup_client(self, config: Config) -> None:
 
         self.wandb_reporter = ClientWandBReporter.from_config(self.client_name, config)
 
-        self.intialized = True
+        self.initialized = True
 
     def get_parameter_exchanger(self, config: Config) -> ParameterExchanger:
         """

fl4health/clients/tabular_data_client.py

@@ -9,7 +9,7 @@
 from sklearn.pipeline import Pipeline
 
 from fl4health.clients.basic_client import BasicClient
-from fl4health.feature_alignment.constants import FEATURE_INFO, FORMAT_SPECIFIED, INPUT_DIMENSION, OUTPUT_DIMENSION
+from fl4health.feature_alignment.constants import FEATURE_INFO, INPUT_DIMENSION, OUTPUT_DIMENSION, SOURCE_SPECIFIED
 from fl4health.feature_alignment.tab_features_info_encoder import TabularFeaturesInfoEncoder
 from fl4health.feature_alignment.tab_features_preprocessor import TabularFeaturesPreprocessor
 from fl4health.utils.metrics import Metric
@@ -43,17 +43,17 @@ def setup_client(self, config: Config) -> None:
         Initialize the client by encoding the information of its tabular data
         and initializing the corresponding TabularFeaturesPreprocessor.
 
-        config[FORMAT_SPECIFIED] indicates whether the server has obtained
+        config[SOURCE_SPECIFIED] indicates whether the server has obtained
         the source of information to perform feature alignment.
         If it is True, it means the server has obtained such information
         (either a priori or by polling a client).
         So the client will encode that information and use it instead
         to perform feature preprocessing.
         """
-        format_specified = self.narrow_config_type(config, FORMAT_SPECIFIED, bool)
+        source_specified = self.narrow_config_type(config, SOURCE_SPECIFIED, bool)
         self.df = self.get_data_frame(config)
 
-        if format_specified:
+        if source_specified:
             # Since the server has obtained its source of information,
             # the client will encode that instead.
             self.tabular_features_info_encoder = TabularFeaturesInfoEncoder.from_json(
@@ -74,7 +74,7 @@ def setup_client(self, config: Config) -> None:
             # that the first dimension is the number of rows.
             self.input_dimension = self.aligned_features.shape[1]
             self.output_dimension = self.tabular_features_info_encoder.get_target_dimension()
-            log(INFO, f"input dimension: {self.input_dimension}, output_dimension: {self.output_dimension}")
+            log(INFO, f"input dimension: {self.input_dimension}, target dimension: {self.output_dimension}")
 
             super().setup_client(config)
 
@@ -111,8 +111,8 @@ def get_properties(self, config: Config) -> Dict[str, Scalar]:
         """
         if not self.initialized:
             self.setup_client(config)
-        format_specified = self.narrow_config_type(config, FORMAT_SPECIFIED, bool)
-        if not format_specified:
+        source_specified = self.narrow_config_type(config, SOURCE_SPECIFIED, bool)
+        if not source_specified:
             return {
                 FEATURE_INFO: self.tabular_features_info_encoder.to_json(),
             }

fl4health/feature_alignment/constants.py

@@ -9,7 +9,7 @@
 
 # FORMAT_SPECIFIED indicates whether the server has the "source of truth"
 # to be used for feature alignment.
-FORMAT_SPECIFIED = "format_specified"
+SOURCE_SPECIFIED = "source_specified"
 
 # FEATURE_INFO refers to the encoded feature information (source of truth).
 FEATURE_INFO = "feature_info"

fl4health/server/tabular_feature_alignment_server.py

@@ -13,9 +13,9 @@
 from fl4health.feature_alignment.constants import (
     CURRENT_SERVER_ROUND,
     FEATURE_INFO,
-    FORMAT_SPECIFIED,
     INPUT_DIMENSION,
     OUTPUT_DIMENSION,
+    SOURCE_SPECIFIED,
 )
 from fl4health.feature_alignment.tab_features_info_encoder import TabularFeaturesInfoEncoder
 from fl4health.reporting.fl_wanb import ServerWandBReporter
@@ -70,11 +70,11 @@ def __init__(
         self.tab_features_info = tabular_features_source_of_truth
         self.config = config
         self.initialize_parameters = initialize_parameters
-        self.format_info_gathered = False
+        self.source_info_gathered = False
         self.dimension_info: Dict[str, int] = {}
         # ensure that self.strategy has type BasicFedAvg so its on_fit_config_fn can be specified.
         assert isinstance(self.strategy, BasicFedAvg)
-        self.strategy.on_fit_config_fn = partial(fit_config, self.config, self.format_info_gathered)
+        self.strategy.on_fit_config_fn = partial(fit_config, self.config, self.source_info_gathered)
 
     def _set_dimension_info(self, input_dimension: int, output_dimension: int) -> None:
         self.dimension_info[INPUT_DIMENSION] = input_dimension
@@ -110,9 +110,9 @@ def fit(self, num_rounds: int, timeout: Optional[float]) -> History:
 
             # the feature information is sent to clients through the config parameter.
             self.config[FEATURE_INFO] = feature_info_source
-            self.format_info_gathered = True
+            self.source_info_gathered = True
 
-            self.strategy.on_fit_config_fn = partial(fit_config, self.config, self.format_info_gathered)
+            self.strategy.on_fit_config_fn = partial(fit_config, self.config, self.source_info_gathered)
 
             # Now the server waits until feature alignment is performed on the clients' side
             # and subsequently requests the input and output dimensions, which are needed for initializing
@@ -158,7 +158,7 @@ def poll_clients_for_dimension_info(self, timeout: Optional[float]) -> Tuple[int
         return input_dimension, output_dimension
 
 
-def fit_config(config: Config, format_specified: bool, current_server_round: int) -> Config:
-    config[FORMAT_SPECIFIED] = format_specified
+def fit_config(config: Config, source_specified: bool, current_server_round: int) -> Config:
+    config[SOURCE_SPECIFIED] = source_specified
     config[CURRENT_SERVER_ROUND] = current_server_round
     return config

tests/clients/test_setup_client_basic.py

@@ -0,0 +1,44 @@
+from pathlib import Path
+from typing import Tuple
+
+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.basic_client import BasicClient
+from fl4health.utils.metrics import Accuracy
+from tests.test_utils.models_for_test import LinearModel
+
+
+class TestingClient(BasicClient):
+    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 LinearModel().to(self.device)
+
+
+def test_setup_client() -> None:
+    client = TestingClient(data_path=Path(""), metrics=[Accuracy()], device=torch.device("cpu"))
+    client.setup_client({})
+    assert client.parameter_exchanger is not None
+    assert client.model is not None
+    assert client.optimizer is not None
+    assert client.train_loader is not None
+    assert client.val_loader is not None
+    assert client.num_train_samples is not None
+    assert client.num_val_samples is not None
+    assert client.learning_rate is not None
+    assert client.initialized