Adding checkpointer and metrics, re-shuffling data in labels split

declearn/quickrun/_config.py

@@ -18,7 +18,7 @@
 """TOML-parsable container for quickrun configurations."""
 
 import dataclasses
-from typing import Dict, List, Literal, Optional, Union
+from typing import Dict, List, Optional, Union
 
 from declearn.utils import TomlConfig
 
@@ -40,11 +40,13 @@ class ModelConfig(TomlConfig):
 
 @dataclasses.dataclass
 class DataSplitConfig(TomlConfig):
-    """Dataclass associated with the function
-    declearn.quickrun._split_data:split_data
+    """Dataclass associated with the functions
+    declearn.quickrun._split_data:split_data and
+    declearn.quickrun._parser:parse_data_folder
 
-    export_folder: str
-        Path to the folder where to export shard-wise files.
+    data_folder: str
+        Absolute path to the folder where to export shard-wise files,
+        and/or to the main folder hosting the data.
     n_shards: int
         Number of shards between which to split the data.
     data_file: str or None, default=None
@@ -67,37 +69,44 @@ class DataSplitConfig(TomlConfig):
         ]0,1] range.
     seed: int or None, default=None
         Optional seed to the RNG used for all sampling operations.
+    client_names: list or None
+        List of custom client names to look for in the data_folder.
+        If None, default to expected prefix search.
+    dataset_names: dict or None
+        Dict of custom dataset names, to look for in each client folder.
+        Expect 'train_data, train_target, valid_data, valid_target' as keys.
+        If None, , default to expected prefix search.
     """
 
-    export_folder: str = "."
+    # Common args
+    data_folder: Optional[str] = None
+    # split_data args
     n_shards: int = 5
     data_file: Optional[str] = None
     label_file: Optional[Union[str, int]] = None
-    scheme: Literal["iid", "labels", "biased"] = "iid"
+    scheme: str = "iid"
     perc_train: float = 0.8
     seed: Optional[int] = None
+    # parse_data_folder args
+    client_names: Optional[List[str]] = None
+    dataset_names: Optional[Dict[str, str]] = None
 
 
 @dataclasses.dataclass
 class ExperimentConfig(TomlConfig):
     """
 
-    Dataclass associated with the function
-    declearn.quickrun._parser:parse_data_folder
+    Dataclass providing kwargs to
+    declearn.main._server.FederatedServer
+    and declearn.main._client.FederatedClient
 
-    data_folder : str or none
-        Absolute path to the main folder hosting the data, overwriting
-        the folder argument if provided. If None, default to expected
-        prefix search in folder.
-    client_names: list or None
-        List of custom client names to look for in the data_folder.
-        If None, default to expected prefix search.
-    dataset_names: dict or None
-        Dict of custom dataset names, to look for in each client folder.
-        Expect 'train_data, train_target, valid_data, valid_target' as keys.
-        If None, , default to expected prefix search.
+    metrics: list[str] or None
+        List of Metric childclass names, defining evaluation metrics
+        to compute in addition to the model's loss.
+    checkpoint: str or None
+        The checkpoint folder path and use default values for other parameters
+        to be used so as to save round-wise model
     """
 
-    data_folder: Optional[str] = None
-    client_names: Optional[List[str]] = None
-    dataset_names: Optional[Dict[str, str]] = None
+    metrics: Optional[List[str]] = None
+    checkpoint: Optional[str] = None

declearn/quickrun/_split_data.py

@@ -167,7 +167,8 @@ def _split_labels(
         srt = idx * s_len
         end = (srt + s_len) if idx < (n_shards - 1) else len(order)
         shard = np.isin(target, order[srt:end])
-        split.append((inputs[shard], target[shard]))
+        shuffle = rng.permutation(shard.sum())
+        split.append((inputs[shard][shuffle], target[shard][shuffle]))
     return split
 
 
@@ -198,7 +199,7 @@ def _split_biased(
     return split
 
 
-def split_data(data_config: DataSplitConfig) -> None:
+def split_data(data_config: DataSplitConfig, folder: str) -> None:
     """Download and randomly split a dataset into shards.
 
     The resulting folder structure is :
@@ -223,6 +224,9 @@ def split_data(data_config: DataSplitConfig) -> None:
         os.makedirs(data_dir, exist_ok=True)
         np.save(os.path.join(data_dir, f"{name}.npy"), data)
 
+    # Overwrite default folder if provided
+    if data_config.data_folder:
+        folder = data_config.data_folder
     # Select the splitting function to be used.
     scheme = data_config.scheme
     if scheme == "iid":
@@ -242,7 +246,7 @@ def split_data(data_config: DataSplitConfig) -> None:
     )
     split = func(inputs, labels, data_config.n_shards, rng)
     # Export the resulting shard-wise data to files.
-    folder = os.path.join(data_config.export_folder, f"data_{scheme}")
+    folder = os.path.join(folder, f"data_{scheme}")
     for i, (inp, tgt) in enumerate(split):
         perc_train = data_config.perc_train
         if not perc_train:

declearn/quickrun/run.py

@@ -36,7 +36,6 @@ import importlib
 import os
 import re
 import textwrap
-from pathlib import Path
 from typing import Any, Dict, List, Optional, Tuple
 
 from declearn.communication import NetworkClientConfig, NetworkServerConfig
@@ -77,10 +76,18 @@ def _run_server(
     model_config: ModelConfig,
     optim: FLOptimConfig,
     config: FLRunConfig,
+    expe_config: ExperimentConfig,
 ) -> None:
     """Routine to run a FL server, called by `run_declearn_experiment`."""
     model = _get_model(folder, model_config)
-    server = FederatedServer(model, network, optim)
+    if expe_config.checkpoint:
+        checkpoint = expe_config.checkpoint
+    else:
+        checkpoint = os.path.join(folder, "result")
+    checkpoint = os.path.join(checkpoint, "server")
+    server = FederatedServer(
+        model, network, optim, expe_config.metrics, checkpoint
+    )
     server.run(config)
 
 
@@ -88,6 +95,7 @@ def _run_client(
     folder: str,
     network: NetworkClientConfig,
     model_config: ModelConfig,
+    expe_config: ExperimentConfig,
     name: str,
     paths: dict,
 ) -> None:
@@ -96,6 +104,12 @@ def _run_client(
     network.name = name
     # Make the model importable
     _ = _get_model(folder, model_config)
+    # Add checkpointer
+    if expe_config.checkpoint:
+        checkpoint = expe_config.checkpoint
+    else:
+        checkpoint = os.path.join(folder, "result")
+    checkpoint = os.path.join(checkpoint, name)
     # Wrap train and validation data as Dataset objects.
     train = InMemoryDataset(
         paths.get("train_data"),
@@ -106,7 +120,7 @@ def _run_client(
         paths.get("valid_data"),
         target=paths.get("valid_target"),
     )
-    client = FederatedClient(network, train, valid)
+    client = FederatedClient(network, train, valid, checkpoint)
     client.run()
 
 
@@ -132,22 +146,29 @@ def get_toml_folder(config: Optional[str] = None) -> Tuple[str, str]:
 
 def locate_or_create_split_data(toml: str, folder: str) -> Dict:
     """Attempts to find split data according to the config toml or
-    or the defualt behavior. If failed, attempts to find full data
+    or the default behavior. If failed, attempts to find full data
     according to the config toml and split it"""
-    expe_config = ExperimentConfig.from_toml(toml, False, "experiment")
+    data_config = DataSplitConfig.from_toml(toml, False, "data")
     try:
-        client_dict = parse_data_folder(expe_config, folder)
+        client_dict = parse_data_folder(data_config, folder)
     except ValueError:
-        data_config = DataSplitConfig.from_toml(toml, False, "data")
-        split_data(folder, data_config)
-        client_dict = parse_data_folder(expe_config,folder)
+        split_data(data_config, folder)
+        client_dict = parse_data_folder(data_config, folder)
     return client_dict
 
 
-def quickrun(
-    config: Optional[str] = None,
-    **kwargs: Any,
-) -> None:
+def server_to_client_network(
+    network_cfg: NetworkServerConfig,
+) -> NetworkClientConfig:
+    "Converts server network config to client network config"
+    return NetworkClientConfig.from_params(
+        protocol=network_cfg.protocol,
+        server_uri=f"ws://localhost:{network_cfg.port}",
+        name="replaceme",
+    )
+
+
+def quickrun(config: Optional[str] = None) -> None:
     """Run a server and its clients using multiprocessing.
 
     The kwargs are the arguments expected by split_data,
@@ -156,20 +177,24 @@ def quickrun(
     toml, folder = get_toml_folder(config)
     # locate split data or split it if needed
     client_dict = locate_or_create_split_data(toml, folder)
-    # Parse toml file to ServerConfig and ClientConfig
-    ntk_server = NetworkServerConfig.from_toml(toml, False, "network_server")
-    optim = FLOptimConfig.from_toml(toml, False, "optim")
-    run = FLRunConfig.from_toml(toml, False, "run")
-    ntk_client = NetworkClientConfig.from_toml(toml, False, "network_client")
-    model_config = ModelConfig.from_toml(toml, False, "model")
+    # Parse toml files
+    ntk_server_cfg = NetworkServerConfig.from_toml(toml, False, "network")
+    ntk_client_cfg = server_to_client_network(ntk_server_cfg)
+    optim_cgf = FLOptimConfig.from_toml(toml, False, "optim")
+    run_cfg = FLRunConfig.from_toml(toml, False, "run")
+    model_cfg = ModelConfig.from_toml(toml, False, "model")
+    expe_cfg = ExperimentConfig.from_toml(toml, False, "experiment")
     # Set up a (func, args) tuple specifying the server process.
-    p_server = (_run_server, (folder, ntk_server, model_config, optim, run))
+    p_server = (
+        _run_server,
+        (folder, ntk_server_cfg, model_cfg, optim_cgf, run_cfg, expe_cfg),
+    )
     # Set up the (func, args) tuples specifying client-wise processes.
     p_client = []
     for name, data_dict in client_dict.items():
         client = (
             _run_client,
-            (folder, ntk_client, model_config, name, data_dict),
+            (folder, ntk_client_cfg, model_cfg, expe_cfg, name, data_dict),
         )
         p_client.append(client)
     # Run each and every process in parallel.
@@ -211,7 +236,7 @@ def parse_args(args: Optional[List[str]] = None) -> argparse.Namespace:
 def main(args: Optional[List[str]] = None) -> None:
     """Quikcrun based on commandline-input arguments."""
     cmdargs = parse_args(args)
-    quickrun(folder=cmdargs.config)
+    quickrun(config=cmdargs.config)
 
 
 if __name__ == "__main__":

examples/quickrun/config.toml

-[network_server]
+[network]
 protocol = "websockets"
 host = "127.0.0.1"
 port = 8765
 
-[network_client]
-protocol = "websockets"
-server_uri = "ws://localhost:8765"
-name = "replaceme"
-
 [optim]
 aggregator = "averaging" # The chosen aggregation strategy
 server_opt = 1.0 # The server learning rate
 
     [optim.client_opt]
-    lrate = 0.01 # The client learning rate
-    regularizers = [["lasso", {alpha = 0.1}]] # The list of regularizer modules, each a list
-    modules = [["momentum", {"beta" = 0.9}]]
+    lrate = 0.001 # The client learning rate
+    modules = ["adam"] # The optimzer modules 
 
 [run]
-rounds = 2 # Number of training rounds
+rounds = 10 # Number of training rounds
 
     [run.register]
     min_clients = 1
@@ -36,10 +30,11 @@ rounds = 2 # Number of training rounds
 
 [experiment]
 # all args for parse_data_folder
-# target_folder = !!!
+
 
 [model]
 # information on where to find the model file
 
 [data]
-# all args from split_data argparser
+# all args from split_data argparser
\ No newline at end of file
+scheme = "labels"
\ No newline at end of file