xgboost-comprehensive with bagging aggregation

examples/quickstart-xgboost/README.md

@@ -0,0 +1,75 @@
+# Flower Example using XGBoost
+
+This example demonstrates how to perform EXtreme Gradient Boosting (XGBoost) within Flower using `xgboost` package.
+Tree-based with bagging method is used for aggregation on the server.
+
+## Project Setup
+
+Start by cloning the example project. We prepared a single-line command that you can copy into your shell which will checkout the example for you:
+
+```shell
+git clone --depth=1 https://github.com/adap/flower.git && mv flower/examples/quickstart-xgboost . && rm -rf flower && cd quickstart-xgboost
+```
+
+This will create a new directory called `quickstart-xgboost` containing the following files:
+
+```shell
+-- pyproject.toml
+-- requirements.txt
+-- client.py
+-- server.py
+-- strategy.py
+-- dataset.py
+-- README.md
+```
+
+### Installing Dependencies
+
+Project dependencies (such as `xgboost` and `flwr`) are defined in `pyproject.toml` and `requirements.txt`. We recommend [Poetry](https://python-poetry.org/docs/) to install those dependencies and manage your virtual environment ([Poetry installation](https://python-poetry.org/docs/#installation)) or [pip](https://pip.pypa.io/en/latest/development/), but feel free to use a different way of installing dependencies and managing virtual environments if you have other preferences.
+
+#### Poetry
+
+```shell
+poetry install
+poetry shell
+```
+
+Poetry will install all your dependencies in a newly created virtual environment. To verify that everything works correctly you can run the following command:
+
+```shell
+poetry run python3 -c "import flwr"
+```
+
+If you don't see any errors you're good to go!
+
+#### pip
+
+Write the command below in your terminal to install the dependencies according to the configuration file requirements.txt.
+
+```shell
+pip install -r requirements.txt
+```
+
+## Run Federated Learning with XGBoost and Flower
+
+Afterwards you are ready to start the Flower server as well as the clients. You can simply start the server in a terminal as follows:
+
+```shell
+python3 server.py
+```
+
+Now you are ready to start the Flower clients which will participate in the learning. To do so simply open two more terminal windows and run the following commands.
+
+Start client 1 in the first terminal:
+
+```shell
+python3 client.py
+```
+
+Start client 2 in the second terminal:
+
+```shell
+python3 client.py
+```
+
+You will see that XGBoost is starting a federated training. Look at the [code](https://github.com/adap/flower/tree/main/examples/quickstart-xgboost) and [tutorial](https://flower.dev/docs/framework/tutorial-quickstart-xgboost.html) for a detailed explanation.

examples/quickstart-xgboost/client.py

@@ -0,0 +1,148 @@
+import warnings
+import xgboost as xgb
+
+import flwr as fl
+from flwr_datasets import FederatedDataset
+from flwr.common import (
+    Code,
+    EvaluateIns,
+    EvaluateRes,
+    FitIns,
+    FitRes,
+    GetParametersIns,
+    GetParametersRes,
+    Parameters,
+    Status,
+)
+
+from dataset import instantiate_partitioner, train_test_split, transform_dataset_to_dmatrix
+
+
+warnings.filterwarnings("ignore", category=UserWarning)
+
+# Load (HIGGS) dataset and conduct partitioning
+num_partitions = 20
+# partitioner type is chosen from ["uniform", "linear", "square", "exponential"]
+partitioner_type = "uniform"
+
+# instantiate partitioner
+partitioner = instantiate_partitioner(partitioner_type=partitioner_type, num_partitions=num_partitions)
+fds = FederatedDataset(dataset="jxie/higgs", partitioners={"train": partitioner})
+
+# let's use the first partition as an example
+partition_id = 0
+partition = fds.load_partition(idx=partition_id, split="train")
+partition.set_format("numpy")
+
+# train/test splitting and data re-formatting
+SEED = 42
+test_fraction = 0.2
+train_data, valid_data, num_train, num_val = train_test_split(partition, test_fraction=test_fraction, seed=SEED)
+
+# reformat data to DMatrix for xgboost
+train_dmatrix = transform_dataset_to_dmatrix(train_data)
+valid_dmatrix = transform_dataset_to_dmatrix(valid_data)
+
+
+# Hyper-parameters for xgboost training
+num_local_round = 1
+params = {
+    "objective": "binary:logistic",
+    "eta": 0.1,  # lr
+    "max_depth": 8,
+    "eval_metric": "auc",
+    "nthread": 16,
+    "num_parallel_tree": 1,
+    "subsample": 1,
+    "tree_method": "hist",
+}
+
+
+# Define Flower client
+class FlowerClient(fl.client.Client):
+    def __init__(self):
+        self.bst = None
+
+    def get_parameters(self, ins: GetParametersIns) -> GetParametersRes:
+        _ = (self, ins)
+        return GetParametersRes(
+            status=Status(
+                code=Code.OK,
+                message="OK",
+            ),
+            parameters=Parameters(tensor_type="", tensors=[]),
+        )
+
+    def local_boost(self):
+        # update trees based on local training data.
+        for i in range(num_local_round):
+            self.bst.update(train_dmatrix, self.bst.num_boosted_rounds())
+
+        # extract the last N=num_local_round trees as new local model
+        bst = self.bst[
+            self.bst.num_boosted_rounds()
+            - num_local_round : self.bst.num_boosted_rounds()
+        ]
+        return bst
+
+    def fit(self, ins: FitIns) -> FitRes:
+        if not self.bst:
+            # first round local training
+            print("Start training at round 1")
+            bst = xgb.train(
+                params,
+                train_dmatrix,
+                num_boost_round=num_local_round,
+                evals=[(valid_dmatrix, "validate"), (train_dmatrix, "train")],
+            )
+            self.config = bst.save_config()
+            self.bst = bst
+        else:
+            print("load global model")
+            for item in ins.parameters.tensors:
+                global_model = bytearray(item)
+
+            # load global model into booster
+            self.bst.load_model(global_model)
+            self.bst.load_config(self.config)
+
+            bst = self.local_boost()
+
+        local_model = bst.save_raw("json")
+        local_model_bytes = bytes(local_model)
+
+        return FitRes(
+            status=Status(
+                code=Code.OK,
+                message="OK",
+            ),
+            parameters=Parameters(tensor_type="", tensors=[local_model_bytes]),
+            num_examples=num_train,
+            metrics={},
+        )
+
+    def evaluate(self, ins: EvaluateIns) -> EvaluateRes:
+        eval_results = self.bst.eval_set(
+            evals=[(train_dmatrix, "train"), (valid_dmatrix, "valid")],
+            iteration=self.bst.num_boosted_rounds() - 1,
+        )
+        auc = round(float(eval_results.split("\t")[2].split(":")[1]), 4)
+
+        global_round = ins.config["global_round"]
+        print(f"AUC = {auc} at round {global_round}")
+
+        return EvaluateRes(
+            status=Status(
+                code=Code.OK,
+                message="OK",
+            ),
+            loss=0.0,
+            num_examples=num_val,
+            metrics={"AUC": auc},
+        )
+
+
+# Start Flower client
+fl.client.start_client(
+    server_address="127.0.0.1:8080", client=FlowerClient().to_client()
+)

examples/quickstart-xgboost/dataset.py

@@ -0,0 +1,42 @@
+import datasets
+import xgboost as xgb
+
+from flwr_datasets.partitioner import (
+    IidPartitioner,
+    LinearPartitioner,
+    SquarePartitioner,
+    ExponentialPartitioner,
+)
+
+CORRELATION_TO_PARTITIONER = {
+    "uniform": IidPartitioner,
+    "linear": LinearPartitioner,
+    "square": SquarePartitioner,
+    "exponential": ExponentialPartitioner,
+}
+
+
+def instantiate_partitioner(partitioner_type: str, num_partitions: int):
+    """Initialise partitioner based on selected partitioner type and number of partitions."""
+    partitioner = CORRELATION_TO_PARTITIONER[partitioner_type](num_partitions=num_partitions)
+    return partitioner
+
+
+def train_test_split(partition: datasets.Dataset, test_fraction: float, seed: int):
+    """Split the data into train and validation set given split rate."""
+    train_test = partition.train_test_split(test_size=test_fraction, seed=seed)
+    partition_train = train_test["train"]
+    partition_test = train_test["test"]
+
+    num_train = len(partition_train)
+    num_test = len(partition_test)
+
+    return partition_train, partition_test, num_train, num_test
+
+
+def transform_dataset_to_dmatrix(data):
+    """transform dataset to DMatrix format for xgboost."""
+    x = data["inputs"]
+    y = data["label"]
+    new_data = xgb.DMatrix(x, label=y)
+    return new_data

examples/quickstart-xgboost/pyproject.toml

@@ -0,0 +1,16 @@
+[build-system]
+requires = ["poetry-core>=1.4.0"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.poetry]
+name = "quickstart-xgboost"
+version = "0.1.0"
+description = "Federated XGBoost Quickstart with Flower"
+authors = ["The Flower Authors <[email protected]>"]
+
+[tool.poetry.dependencies]
+python = ">=3.8,<3.11"
+flwr = ">=1.0,<2.0"
+torch = "1.13.1"
+torchvision = "0.14.1"
+tqdm = "4.65.0"

examples/quickstart-xgboost/requirements.txt

@@ -0,0 +1,4 @@
+flwr>=1.0, <2.0
+flwr_datasets
+xgboost==2.0.1
+tqdm==4.65.0

examples/quickstart-xgboost/run.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+set -e
+cd "$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"/
+
+# Download the CIFAR-10 dataset
+python -c "from dataset import init_higgs; init_higgs(20, "uniform")"
+
+echo "Starting server"
+python server.py &
+sleep 3  # Sleep for 3s to give the server enough time to start
+
+for i in `seq 0 1`; do
+    echo "Starting client $i"
+    python client.py &
+done
+
+# Enable CTRL+C to stop all background processes
+trap "trap - SIGTERM && kill -- -$$" SIGINT SIGTERM
+# Wait for all background processes to complete
+wait

examples/quickstart-xgboost/server.py

@@ -0,0 +1,45 @@
+from typing import Dict
+
+import flwr as fl
+from strategy import XGbBagging
+
+pool_size = 2
+num_rounds = 5
+num_clients_per_round = 2
+min_evaluate_clients = 2
+
+
+def eval_config(rnd: int) -> Dict[str, str]:
+    """Return a configuration with global epochs."""
+    config = {
+        "global_round": str(rnd),
+    }
+    return config
+
+
+def evaluate_metrics_aggregation(eval_metrics):
+    """Return an aggregated metric (AUC) for evaluation."""
+    auc_aggregated = sum([metrics["AUC"] for _, metrics in eval_metrics]) / len(
+        eval_metrics
+    )
+    metrics_aggregated = {"AUC": auc_aggregated}
+    return metrics_aggregated
+
+
+# Define strategy
+strategy = XGbBagging(
+    fraction_fit=(float(num_clients_per_round) / pool_size),
+    min_fit_clients=num_clients_per_round,
+    min_available_clients=pool_size,
+    fraction_evaluate=1.0,
+    min_evaluate_clients=min_evaluate_clients,
+    on_evaluate_config_fn=eval_config,
+    evaluate_metrics_aggregation_fn=evaluate_metrics_aggregation,
+)
+
+# Start Flower server
+fl.server.start_server(
+    server_address="0.0.0.0:8080",
+    config=fl.server.ServerConfig(num_rounds=num_rounds),
+    strategy=strategy,
+)