flwr 1.21.0__py3-none-any.whl → 1.23.0__py3-none-any.whl

flwr/cli/new/templates/app/code/client.huggingface.py.tpl

@@ -1,41 +1,67 @@
 """$project_name: A Flower / $framework_str app."""
 
 import torch
-from flwr.client import ClientApp, NumPyClient
-from flwr.common import Context
+from flwr.app import ArrayRecord, Context, Message, MetricRecord, RecordDict
+from flwr.clientapp import ClientApp
 from transformers import AutoModelForSequenceClassification
 
-from $import_name.task import get_weights, load_data, set_weights, test, train
+from $import_name.task import load_data
+from $import_name.task import test as test_fn
+from $import_name.task import train as train_fn
 
+# Flower ClientApp
+app = ClientApp()
+
+
+@app.train()
+def train(msg: Message, context: Context):
+    """Train the model on local data."""
+
+    # Get this client's dataset partition
+    partition_id = context.node_config["partition-id"]
+    num_partitions = context.node_config["num-partitions"]
+    model_name = context.run_config["model-name"]
+    trainloader, _ = load_data(partition_id, num_partitions, model_name)
+
+    # Load model
+    num_labels = context.run_config["num-labels"]
+    net = AutoModelForSequenceClassification.from_pretrained(
+        model_name, num_labels=num_labels
+    )
 
-# Flower client
-class FlowerClient(NumPyClient):
-    def __init__(self, net, trainloader, testloader, local_epochs):
-        self.net = net
-        self.trainloader = trainloader
-        self.testloader = testloader
-        self.local_epochs = local_epochs
-        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-        self.net.to(self.device)
+    # Initialize it with the received weights
+    net.load_state_dict(msg.content["arrays"].to_torch_state_dict())
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    net.to(device)
 
-    def fit(self, parameters, config):
-        set_weights(self.net, parameters)
-        train(self.net, self.trainloader, epochs=self.local_epochs, device=self.device)
-        return get_weights(self.net), len(self.trainloader), {}
+    # Train the model on local data
+    train_loss = train_fn(
+        net,
+        trainloader,
+        context.run_config["local-steps"],
+        device,
+    )
 
-    def evaluate(self, parameters, config):
-        set_weights(self.net, parameters)
-        loss, accuracy = test(self.net, self.testloader, self.device)
-        return float(loss), len(self.testloader), {"accuracy": accuracy}
+    # Construct and return reply Message
+    model_record = ArrayRecord(net.state_dict())
+    metrics = {
+        "train_loss": train_loss,
+        "num-examples": len(trainloader.dataset),
+    }
+    metric_record = MetricRecord(metrics)
+    content = RecordDict({"arrays": model_record, "metrics": metric_record})
+    return Message(content=content, reply_to=msg)
 
 
-def client_fn(context: Context):
+@app.evaluate()
+def evaluate(msg: Message, context: Context):
+    """Evaluate the model on local data."""
 
     # Get this client's dataset partition
     partition_id = context.node_config["partition-id"]
     num_partitions = context.node_config["num-partitions"]
     model_name = context.run_config["model-name"]
-    trainloader, valloader = load_data(partition_id, num_partitions, model_name)
+    _, valloader = load_data(partition_id, num_partitions, model_name)
 
     # Load model
     num_labels = context.run_config["num-labels"]
@@ -43,13 +69,25 @@ def client_fn(context: Context):
         model_name, num_labels=num_labels
     )
 
-    local_epochs = context.run_config["local-epochs"]
-
-    # Return Client instance
-    return FlowerClient(net, trainloader, valloader, local_epochs).to_client()
+    # Initialize it with the received weights
+    net.load_state_dict(msg.content["arrays"].to_torch_state_dict())
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    net.to(device)
 
+    # Evaluate the model on local data
+    val_loss, val_accuracy = test_fn(
+        net,
+        valloader,
+        device,
+    )
 
-# Flower ClientApp
-app = ClientApp(
-    client_fn,
-)
+    # Construct and return reply Message
+    model_record = ArrayRecord(net.state_dict())
+    metrics = {
+        "val_loss": val_loss,
+        "val_accuracy": val_accuracy,
+        "num-examples": len(valloader.dataset),
+    }
+    metric_record = MetricRecord(metrics)
+    content = RecordDict({"arrays": model_record, "metrics": metric_record})
+    return Message(content=content, reply_to=msg)

flwr/cli/new/templates/app/code/client.jax.py.tpl

@@ -1,50 +1,71 @@
 """$project_name: A Flower / $framework_str app."""
 
 import jax
+from flwr.app import ArrayRecord, Context, Message, MetricRecord, RecordDict
+from flwr.clientapp import ClientApp
 
-from flwr.client import ClientApp, NumPyClient
-from flwr.common import Context
-from $import_name.task import (
-    evaluation,
-    get_params,
-    load_data,
-    load_model,
-    loss_fn,
-    set_params,
-    train,
-)
-
-
-# Define Flower Client and client_fn
-class FlowerClient(NumPyClient):
-    def __init__(self, input_dim):
-        self.train_x, self.train_y, self.test_x, self.test_y = load_data()
-        self.grad_fn = jax.grad(loss_fn)
-        self.params = load_model((input_dim,))
-
-    def fit(self, parameters, config):
-        set_params(self.params, parameters)
-        self.params, loss, num_examples = train(
-            self.params, self.grad_fn, self.train_x, self.train_y
-        )
-        return get_params(self.params), num_examples, {"loss": float(loss)}
-
-    def evaluate(self, parameters, config):
-        set_params(self.params, parameters)
-        loss, num_examples = evaluation(
-            self.params, self.grad_fn, self.test_x, self.test_y
-        )
-        return float(loss), num_examples, {"loss": float(loss)}
-
-
-def client_fn(context: Context):
+from $import_name.task import evaluation as evaluation_fn
+from $import_name.task import get_params, load_data, load_model, loss_fn, set_params
+from $import_name.task import train as train_fn
+
+# Flower ClientApp
+app = ClientApp()
+
+
+@app.train()
+def train(msg: Message, context: Context):
+    """Train the model on local data."""
+
+    # Read from config
     input_dim = context.run_config["input-dim"]
 
-    # Return Client instance
-    return FlowerClient(input_dim).to_client()
+    # Load data and model
+    train_x, train_y, _, _ = load_data()
+    model = load_model((input_dim,))
+    grad_fn = jax.grad(loss_fn)
 
+    # Set model parameters
+    ndarrays = msg.content["arrays"].to_numpy_ndarrays()
+    set_params(model, ndarrays)
 
-# Flower ClientApp
-app = ClientApp(
-    client_fn,
-)
+    # Train the model on local data
+    model, loss, num_examples = train_fn(model, grad_fn, train_x, train_y)
+
+    # Construct and return reply Message
+    model_record = ArrayRecord(get_params(model))
+    metrics = {
+        "train_loss": float(loss),
+        "num-examples": num_examples,
+    }
+    metric_record = MetricRecord(metrics)
+    content = RecordDict({"arrays": model_record, "metrics": metric_record})
+    return Message(content=content, reply_to=msg)
+
+
+@app.evaluate()
+def evaluate(msg: Message, context: Context):
+    """Evaluate the model on local data."""
+
+    # Read from config
+    input_dim = context.run_config["input-dim"]
+
+    # Load data and model
+    _, _, test_x, test_y = load_data()
+    model = load_model((input_dim,))
+    grad_fn = jax.grad(loss_fn)
+
+    # Set model parameters
+    ndarrays = msg.content["arrays"].to_numpy_ndarrays()
+    set_params(model, ndarrays)
+
+    # Evaluate the model on local data
+    loss, num_examples = evaluation_fn(model, grad_fn, test_x, test_y)
+
+    # Construct and return reply Message
+    metrics = {
+        "test_loss": float(loss),
+        "num-examples": num_examples,
+    }
+    metric_record = MetricRecord(metrics)
+    content = RecordDict({"metrics": metric_record})
+    return Message(content=content, reply_to=msg)

flwr/cli/new/templates/app/code/client.mlx.py.tpl

@@ -3,10 +3,9 @@
 import mlx.core as mx
 import mlx.nn as nn
 import mlx.optimizers as optim
+from flwr.app import ArrayRecord, Context, Message, MetricRecord, RecordDict
+from flwr.clientapp import ClientApp
 
-from flwr.client import ClientApp, NumPyClient
-from flwr.common import Context
-from flwr.common.config import UserConfig
 from $import_name.task import (
     MLP,
     batch_iterate,
@@ -17,57 +16,87 @@ from $import_name.task import (
     set_params,
 )
 
+# Flower ClientApp
+app = ClientApp()
+
+
+@app.train()
+def train(msg: Message, context: Context):
+    """Train the model on local data."""
+
+    # Read config
+    num_layers = context.run_config["num-layers"]
+    input_dim = context.run_config["input-dim"]
+    hidden_dim = context.run_config["hidden-dim"]
+    batch_size = context.run_config["batch-size"]
+    learning_rate = context.run_config["lr"]
+    num_epochs = context.run_config["local-epochs"]
 
-# Define Flower Client and client_fn
-class FlowerClient(NumPyClient):
-    def __init__(
-        self,
-        data,
-        run_config: UserConfig,
-        num_classes,
-    ):
-        num_layers = run_config["num-layers"]
-        hidden_dim = run_config["hidden-dim"]
-        input_dim = run_config["input-dim"]
-        batch_size = run_config["batch-size"]
-        learning_rate = run_config["lr"]
-        self.num_epochs = run_config["local-epochs"]
-
-        self.train_images, self.train_labels, self.test_images, self.test_labels = data
-        self.model = MLP(num_layers, input_dim, hidden_dim, num_classes)
-        self.optimizer = optim.SGD(learning_rate=learning_rate)
-        self.loss_and_grad_fn = nn.value_and_grad(self.model, loss_fn)
-        self.batch_size = batch_size
-
-    def fit(self, parameters, config):
-        set_params(self.model, parameters)
-        for _ in range(self.num_epochs):
-            for X, y in batch_iterate(
-                self.batch_size, self.train_images, self.train_labels
-            ):
-                _, grads = self.loss_and_grad_fn(self.model, X, y)
-                self.optimizer.update(self.model, grads)
-                mx.eval(self.model.parameters(), self.optimizer.state)
-        return get_params(self.model), len(self.train_images), {}
-
-    def evaluate(self, parameters, config):
-        set_params(self.model, parameters)
-        accuracy = eval_fn(self.model, self.test_images, self.test_labels)
-        loss = loss_fn(self.model, self.test_images, self.test_labels)
-        return loss.item(), len(self.test_images), {"accuracy": accuracy.item()}
-
-
-def client_fn(context: Context):
+    # Instantiate model and apply global parameters
+    model = MLP(num_layers, input_dim, hidden_dim, output_dim=10)
+    ndarrays = msg.content["arrays"].to_numpy_ndarrays()
+    set_params(model, ndarrays)
+
+    # Define optimizer and loss function
+    optimizer = optim.SGD(learning_rate=learning_rate)
+    loss_and_grad_fn = nn.value_and_grad(model, loss_fn)
+
+    # Load data
     partition_id = context.node_config["partition-id"]
     num_partitions = context.node_config["num-partitions"]
-    data = load_data(partition_id, num_partitions)
-    num_classes = 10
+    train_images, train_labels, _, _ = load_data(partition_id, num_partitions)
 
-    # Return Client instance
-    return FlowerClient(data, context.run_config, num_classes).to_client()
+    # Train the model on local data
+    for _ in range(num_epochs):
+        for X, y in batch_iterate(batch_size, train_images, train_labels):
+            _, grads = loss_and_grad_fn(model, X, y)
+            optimizer.update(model, grads)
+            mx.eval(model.parameters(), optimizer.state)
 
+    # Compute train accuracy and loss
+    accuracy = eval_fn(model, train_images, train_labels)
+    loss = loss_fn(model, train_images, train_labels)
+    # Construct and return reply Message
+    model_record = ArrayRecord(get_params(model))
+    metrics = {
+        "num-examples": len(train_images),
+        "accuracy": float(accuracy.item()),
+        "loss": float(loss.item()),
+    }
+    metric_record = MetricRecord(metrics)
+    content = RecordDict({"arrays": model_record, "metrics": metric_record})
+    return Message(content=content, reply_to=msg)
 
-# Flower ClientApp
-app = ClientApp(
-    client_fn,
-)
+
+@app.evaluate()
+def evaluate(msg: Message, context: Context):
+    """Evaluate the model on local data."""
+
+    # Read config
+    num_layers = context.run_config["num-layers"]
+    input_dim = context.run_config["input-dim"]
+    hidden_dim = context.run_config["hidden-dim"]
+
+    # Instantiate model and apply global parameters
+    model = MLP(num_layers, input_dim, hidden_dim, output_dim=10)
+    ndarrays = msg.content["arrays"].to_numpy_ndarrays()
+    set_params(model, ndarrays)
+
+    # Load data
+    partition_id = context.node_config["partition-id"]
+    num_partitions = context.node_config["num-partitions"]
+    _, _, test_images, test_labels = load_data(partition_id, num_partitions)
+
+    # Evaluate the model on local data
+    accuracy = eval_fn(model, test_images, test_labels)
+    loss = loss_fn(model, test_images, test_labels)
+
+    # Construct and return reply Message
+    metrics = {
+        "num-examples": len(test_images),
+        "accuracy": float(accuracy.item()),
+        "loss": float(loss.item()),
+    }
+    metric_record = MetricRecord(metrics)
+    content = RecordDict({"metrics": metric_record})
+    return Message(content=content, reply_to=msg)

flwr/cli/new/templates/app/code/client.numpy.py.tpl

@@ -1,23 +1,46 @@
 """$project_name: A Flower / $framework_str app."""
 
-from flwr.client import ClientApp, NumPyClient
-from flwr.common import Context
-from $import_name.task import get_dummy_model
+import numpy as np
+from flwr.app import ArrayRecord, Context, Message, MetricRecord, RecordDict
+from flwr.clientapp import ClientApp
 
+# Flower ClientApp
+app = ClientApp()
 
-class FlowerClient(NumPyClient):
 
-    def fit(self, parameters, config):
-        model = get_dummy_model()
-        return [model], 1, {}
+@app.train()
+def train(msg: Message, context: Context):
+    """Train the model on local data."""
 
-    def evaluate(self, parameters, config):
-        return float(0.0), 1, {"accuracy": float(1.0)}
+    # The model is the global arrays
+    ndarrays = msg.content["arrays"].to_numpy_ndarrays()
 
+    # Simulate local training (here we just add random noise to model parameters)
+    model = [m + np.random.rand(*m.shape) for m in ndarrays]
 
-def client_fn(context: Context):
-    return FlowerClient().to_client()
+    # Construct and return reply Message
+    model_record = ArrayRecord(model)
+    metrics = {
+        "random_metric": np.random.rand(),
+        "num-examples": 1,
+    }
+    metric_record = MetricRecord(metrics)
+    content = RecordDict({"arrays": model_record, "metrics": metric_record})
+    return Message(content=content, reply_to=msg)
 
 
-# Flower ClientApp
-app = ClientApp(client_fn=client_fn)
+@app.evaluate()
+def evaluate(msg: Message, context: Context):
+    """Evaluate the model on local data."""
+
+    # The model is the global arrays
+    ndarrays = msg.content["arrays"].to_numpy_ndarrays()
+
+    # Return reply Message
+    metrics = {
+        "random_metric": np.random.rand(3).tolist(),
+        "num-examples": 1,
+    }
+    metric_record = MetricRecord(metrics)
+    content = RecordDict({"metrics": metric_record})
+    return Message(content=content, reply_to=msg)

flwr/cli/new/templates/app/code/client.pytorch.py.tpl

@@ -1,55 +1,80 @@
 """$project_name: A Flower / $framework_str app."""
 
 import torch
+from flwr.app import ArrayRecord, Context, Message, MetricRecord, RecordDict
+from flwr.clientapp import ClientApp
 
-from flwr.client import ClientApp, NumPyClient
-from flwr.common import Context
-from $import_name.task import Net, get_weights, load_data, set_weights, test, train
-
-
-# Define Flower Client and client_fn
-class FlowerClient(NumPyClient):
-    def __init__(self, net, trainloader, valloader, local_epochs):
-        self.net = net
-        self.trainloader = trainloader
-        self.valloader = valloader
-        self.local_epochs = local_epochs
-        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-        self.net.to(self.device)
-
-    def fit(self, parameters, config):
-        set_weights(self.net, parameters)
-        train_loss = train(
-            self.net,
-            self.trainloader,
-            self.local_epochs,
-            self.device,
-        )
-        return (
-            get_weights(self.net),
-            len(self.trainloader.dataset),
-            {"train_loss": train_loss},
-        )
-
-    def evaluate(self, parameters, config):
-        set_weights(self.net, parameters)
-        loss, accuracy = test(self.net, self.valloader, self.device)
-        return loss, len(self.valloader.dataset), {"accuracy": accuracy}
-
-
-def client_fn(context: Context):
-    # Load model and data
-    net = Net()
+from $import_name.task import Net, load_data
+from $import_name.task import test as test_fn
+from $import_name.task import train as train_fn
+
+# Flower ClientApp
+app = ClientApp()
+
+
+@app.train()
+def train(msg: Message, context: Context):
+    """Train the model on local data."""
+
+    # Load the model and initialize it with the received weights
+    model = Net()
+    model.load_state_dict(msg.content["arrays"].to_torch_state_dict())
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+
+    # Load the data
     partition_id = context.node_config["partition-id"]
     num_partitions = context.node_config["num-partitions"]
-    trainloader, valloader = load_data(partition_id, num_partitions)
-    local_epochs = context.run_config["local-epochs"]
+    trainloader, _ = load_data(partition_id, num_partitions)
 
-    # Return Client instance
-    return FlowerClient(net, trainloader, valloader, local_epochs).to_client()
+    # Call the training function
+    train_loss = train_fn(
+        model,
+        trainloader,
+        context.run_config["local-epochs"],
+        msg.content["config"]["lr"],
+        device,
+    )
 
+    # Construct and return reply Message
+    model_record = ArrayRecord(model.state_dict())
+    metrics = {
+        "train_loss": train_loss,
+        "num-examples": len(trainloader.dataset),
+    }
+    metric_record = MetricRecord(metrics)
+    content = RecordDict({"arrays": model_record, "metrics": metric_record})
+    return Message(content=content, reply_to=msg)
 
-# Flower ClientApp
-app = ClientApp(
-    client_fn,
-)
+
+@app.evaluate()
+def evaluate(msg: Message, context: Context):
+    """Evaluate the model on local data."""
+
+    # Load the model and initialize it with the received weights
+    model = Net()
+    model.load_state_dict(msg.content["arrays"].to_torch_state_dict())
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+
+    # Load the data
+    partition_id = context.node_config["partition-id"]
+    num_partitions = context.node_config["num-partitions"]
+    _, valloader = load_data(partition_id, num_partitions)
+
+    # Call the evaluation function
+    eval_loss, eval_acc = test_fn(
+        model,
+        valloader,
+        device,
+    )
+
+    # Construct and return reply Message
+    metrics = {
+        "eval_loss": eval_loss,
+        "eval_acc": eval_acc,
+        "num-examples": len(valloader.dataset),
+    }
+    metric_record = MetricRecord(metrics)
+    content = RecordDict({"metrics": metric_record})
+    return Message(content=content, reply_to=msg)

flwr/cli/new/templates/app/code/client.pytorch_legacy_api.py.tpl

@@ -0,0 +1,55 @@
+"""$project_name: A Flower / $framework_str app."""
+
+import torch
+
+from flwr.client import ClientApp, NumPyClient
+from flwr.common import Context
+from $import_name.task import Net, get_weights, load_data, set_weights, test, train
+
+
+# Define Flower Client and client_fn
+class FlowerClient(NumPyClient):
+    def __init__(self, net, trainloader, valloader, local_epochs):
+        self.net = net
+        self.trainloader = trainloader
+        self.valloader = valloader
+        self.local_epochs = local_epochs
+        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        self.net.to(self.device)
+
+    def fit(self, parameters, config):
+        set_weights(self.net, parameters)
+        train_loss = train(
+            self.net,
+            self.trainloader,
+            self.local_epochs,
+            self.device,
+        )
+        return (
+            get_weights(self.net),
+            len(self.trainloader.dataset),
+            {"train_loss": train_loss},
+        )
+
+    def evaluate(self, parameters, config):
+        set_weights(self.net, parameters)
+        loss, accuracy = test(self.net, self.valloader, self.device)
+        return loss, len(self.valloader.dataset), {"accuracy": accuracy}
+
+
+def client_fn(context: Context):
+    # Load model and data
+    net = Net()
+    partition_id = context.node_config["partition-id"]
+    num_partitions = context.node_config["num-partitions"]
+    trainloader, valloader = load_data(partition_id, num_partitions)
+    local_epochs = context.run_config["local-epochs"]
+
+    # Return Client instance
+    return FlowerClient(net, trainloader, valloader, local_epochs).to_client()
+
+
+# Flower ClientApp
+app = ClientApp(
+    client_fn,
+)