Update CLI new command templates

examples/app-pytorch/task.py

@@ -35,6 +35,14 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.fc3(x)
 
 
+def load_data():
+    """Load CIFAR-10 (training and test set)."""
+    trf = Compose([ToTensor(), Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
+    trainset = CIFAR10("./data", train=True, download=True, transform=trf)
+    testset = CIFAR10("./data", train=False, download=True, transform=trf)
+    return DataLoader(trainset, batch_size=32, shuffle=True), DataLoader(testset)
+
+
 def train(net, trainloader, valloader, epochs, device):
     """Train the model on the training set."""
     print("Starting training...")
@@ -77,14 +85,6 @@ def test(net, testloader):
     return loss, accuracy
 
 
-def load_data():
-    """Load CIFAR-10 (training and test set)."""
-    trf = Compose([ToTensor(), Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
-    trainset = CIFAR10("./data", train=True, download=True, transform=trf)
-    testset = CIFAR10("./data", train=False, download=True, transform=trf)
-    return DataLoader(trainset, batch_size=32, shuffle=True), DataLoader(testset)
-
-
 def get_parameters(net):
     return [val.cpu().numpy() for _, val in net.state_dict().items()]
 

examples/custom-mods/client.py

@@ -58,9 +58,8 @@ def client_fn(cid: str):
 
 def get_wandb_mod(name: str) -> Mod:
     def wandb_mod(msg: Message, context: Context, app: ClientAppCallable) -> Message:
-        """Flower Mod that logs the metrics dictionary returned by the client's
-        fit function to Weights & Biases.
-        """
+        """Flower Mod that logs the metrics dictionary returned by the client's fit
+        function to Weights & Biases."""
         server_round = int(msg.metadata.group_id)
 
         if server_round == 1 and msg.metadata.message_type == MessageType.TRAIN:
@@ -107,9 +106,8 @@ def get_tensorboard_mod(logdir) -> Mod:
     def tensorboard_mod(
         msg: Message, context: Context, app: ClientAppCallable
     ) -> Message:
-        """Flower Mod that logs the metrics dictionary returned by the client's
-        fit function to TensorBoard.
-        """
+        """Flower Mod that logs the metrics dictionary returned by the client's fit
+        function to TensorBoard."""
         logdir_run = os.path.join(logdir, str(msg.metadata.run_id))
 
         node_id = str(msg.metadata.dst_node_id)

src/py/flwr/cli/new/new.py

@@ -97,28 +97,28 @@ def new(
         ]
         framework_str = selected_value[0]
 
+    framework_str = framework_str.lower()
+
     # Set project directory path
     cwd = os.getcwd()
     pnl = project_name.lower()
     project_dir = os.path.join(cwd, pnl)
 
     # List of files to render
     files = {
-        "README.md": {
-            "template": "app/README.md.tpl",
-        },
-        "requirements.txt": {
-            "template": f"app/requirements.{framework_str.lower()}.txt.tpl"
-        },
+        "README.md": {"template": "app/README.md.tpl"},
+        "requirements.txt": {"template": f"app/requirements.{framework_str}.txt.tpl"},
         "flower.toml": {"template": "app/flower.toml.tpl"},
+        "pyproject.toml": {"template": "app/pyproject.toml.tpl"},
         f"{pnl}/__init__.py": {"template": "app/code/__init__.py.tpl"},
-        f"{pnl}/server.py": {
-            "template": f"app/code/server.{framework_str.lower()}.py.tpl"
-        },
-        f"{pnl}/client.py": {
-            "template": f"app/code/client.{framework_str.lower()}.py.tpl"
-        },
+        f"{pnl}/server.py": {"template": f"app/code/server.{framework_str}.py.tpl"},
+        f"{pnl}/client.py": {"template": f"app/code/client.{framework_str}.py.tpl"},
     }
+
+    # In case framework is MlFramework.PYTORCH generate additionally the task.py file
+    if framework_str == MlFramework.PYTORCH.value.lower():
+        files[f"{pnl}/task.py"] = {"template": f"app/code/task.{framework_str}.py.tpl"}
+
     context = {"project_name": project_name}
 
     for file_path, value in files.items():

src/py/flwr/cli/new/new_test.py

@@ -70,11 +70,13 @@ def test_new(tmp_path: str) -> None:
         "fedgpt",
         "README.md",
         "flower.toml",
+        "pyproject.toml",
     }
     expected_files_module = {
         "__init__.py",
         "server.py",
         "client.py",
+        "task.py",
     }
 
     # Current directory

src/py/flwr/cli/new/templates/app/README.md.tpl

@@ -6,13 +6,23 @@
 pip install -r requirements.txt
 ```
 
-## Start the SuperLink
+## Run (Simulation Engine)
+
+In the `$project_name` directory, use `flwr run` to run a local simulation:
+
+```bash
+flwr run
+```
+
+## Run (Deployment Engine)
+
+### Start the SuperLink
 
 ```bash
 flower-superlink --insecure
 ```
 
-## Start the long-running Flower client
+### Start the long-running Flower client
 
 In a new terminal window, start the first long-running Flower client:
 
@@ -26,7 +36,7 @@ In yet another new terminal window, start the second long-running Flower client:
 flower-client-app client:app --insecure
 ```
 
-## Start the ServerApp
+### Start the ServerApp
 
 ```bash
 flower-server-app server:app --insecure

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

@@ -1,11 +1,10 @@
 """$project_name: A Flower / NumPy app."""
 
-import flwr as fl
+from flwr.client import NumPyClient, ClientApp
 import numpy as np
 
 
-# Flower client, adapted from Pytorch quickstart example
-class FlowerClient(fl.client.NumPyClient):
+class FlowerClient(NumPyClient):
     def get_parameters(self, config):
         return [np.ones((1, 1))]
 
@@ -20,5 +19,5 @@ def client_fn(cid: str):
     return FlowerClient().to_client()
 
 
-# ClientApp for Flower-Next
-app = fl.client.ClientApp(client_fn=client_fn)
+# Flower ClientApp
+app = ClientApp(client_fn=client_fn)

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

@@ -1 +1,44 @@
 """$project_name: A Flower / PyTorch app."""
+
+from flwr.client import NumPyClient, ClientApp
+
+from $project_name.task import (
+    Net,
+    DEVICE,
+    load_data,
+    get_weights,
+    set_weights,
+    train,
+    test,
+)
+
+
+# Define Flower Client and client_fn
+class FlowerClient(NumPyClient):
+    def __init__(self, net, trainloader, valloader) -> None:
+        self.net = net
+        self.trainloader = trainloader
+        self.valloader = valloader
+
+    def fit(self, parameters, config):
+        set_weights(self.net, parameters)
+        results = train(self.net, self.trainloader, self.valloader, 1, DEVICE)
+        return get_weights(self.net), len(self.trainloader.dataset), results
+
+    def evaluate(self, parameters, config):
+        set_weights(self.net, parameters)
+        loss, accuracy = test(self.net, self.valloader)
+        return loss, len(self.valloader.dataset), {"accuracy": accuracy}
+
+
+def client_fn(cid: str):
+    # Load model and data
+    net = Net().to(DEVICE)
+    trainloader, valloader = load_data()
+
+    # Return Client instance
+    return FlowerClient(net, trainloader, valloader).to_client()
+
+
+# Flower ClientApp
+app = ClientApp(client_fn)

src/py/flwr/cli/new/templates/app/code/server.pytorch.py.tpl

@@ -1 +1,28 @@
 """$project_name: A Flower / PyTorch app."""
+
+from flwr.common import ndarrays_to_parameters
+from flwr.server import ServerApp, ServerConfig
+from flwr.server.strategy import FedAvg
+
+from $project_name.task import Net, get_weights
+
+
+# Initialize model parameters
+ndarrays = get_weights(Net())
+parameters = ndarrays_to_parameters(ndarrays)
+
+
+# Define strategy
+strategy = FedAvg(
+    fraction_fit=1.0,
+    fraction_evaluate=1.0,
+    min_available_clients=2,
+    initial_parameters=parameters,
+)
+
+
+# Create ServerApp
+app = ServerApp(
+    config=ServerConfig(num_rounds=3),
+    strategy=strategy,
+)

src/py/flwr/cli/new/templates/app/code/task.pytorch.py.tpl

@@ -0,0 +1,96 @@
+"""$project_name: A Flower / PyTorch app."""
+
+import warnings
+from collections import OrderedDict
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from torchvision.datasets import CIFAR10
+from torchvision.transforms import Compose, Normalize, ToTensor
+from tqdm import tqdm
+
+
+DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+
+class Net(nn.Module):
+    """Model (simple CNN adapted from 'PyTorch: A 60 Minute Blitz')"""
+
+    def __init__(self) -> None:
+        super(Net, self).__init__()
+        self.conv1 = nn.Conv2d(3, 6, 5)
+        self.pool = nn.MaxPool2d(2, 2)
+        self.conv2 = nn.Conv2d(6, 16, 5)
+        self.fc1 = nn.Linear(16 * 5 * 5, 120)
+        self.fc2 = nn.Linear(120, 84)
+        self.fc3 = nn.Linear(84, 10)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.pool(F.relu(self.conv1(x)))
+        x = self.pool(F.relu(self.conv2(x)))
+        x = x.view(-1, 16 * 5 * 5)
+        x = F.relu(self.fc1(x))
+        x = F.relu(self.fc2(x))
+        return self.fc3(x)
+
+
+def load_data():
+    """Load CIFAR-10 (training and test set)."""
+    trf = Compose([ToTensor(), Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
+    trainset = CIFAR10("./data", train=True, download=True, transform=trf)
+    testset = CIFAR10("./data", train=False, download=True, transform=trf)
+    return DataLoader(trainset, batch_size=32, shuffle=True), DataLoader(testset)
+
+
+def train(net, trainloader, valloader, epochs, device):
+    """Train the model on the training set."""
+    print("Starting training...")
+    net.to(device)  # move model to GPU if available
+    criterion = torch.nn.CrossEntropyLoss().to(device)
+    optimizer = torch.optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
+    net.train()
+    for _ in range(epochs):
+        for images, labels in trainloader:
+            images, labels = images.to(device), labels.to(device)
+            optimizer.zero_grad()
+            loss = criterion(net(images), labels)
+            loss.backward()
+            optimizer.step()
+
+    train_loss, train_acc = test(net, trainloader)
+    val_loss, val_acc = test(net, valloader)
+
+    results = {
+        "train_loss": train_loss,
+        "train_accuracy": train_acc,
+        "val_loss": val_loss,
+        "val_accuracy": val_acc,
+    }
+    return results
+
+
+def test(net, testloader):
+    """Validate the model on the test set."""
+    net.to(DEVICE)
+    criterion = torch.nn.CrossEntropyLoss()
+    correct, loss = 0, 0.0
+    with torch.no_grad():
+        for images, labels in tqdm(testloader):
+            outputs = net(images.to(DEVICE))
+            labels = labels.to(DEVICE)
+            loss += criterion(outputs, labels).item()
+            correct += (torch.max(outputs.data, 1)[1] == labels).sum().item()
+    accuracy = correct / len(testloader.dataset)
+    return loss, accuracy
+
+
+def get_weights(net):
+    return [val.cpu().numpy() for _, val in net.state_dict().items()]
+
+
+def set_weights(net, parameters):
+    params_dict = zip(net.state_dict().keys(), parameters)
+    state_dict = OrderedDict({k: torch.tensor(v) for k, v in params_dict})
+    net.load_state_dict(state_dict, strict=True)

src/py/flwr/cli/new/templates/app/flower.toml.tpl

@@ -3,7 +3,10 @@ name = "$project_name"
 version = "1.0.0"
 description = ""
 license = "Apache-2.0"
-authors = ["The Flower Authors <[email protected]>"]
+authors = [
+    "The Flower Authors <[email protected]>",
+]
+readme = "README.md"
 
 [flower.components]
 serverapp = "$project_name.server:app"

src/py/flwr/cli/new/templates/app/pyproject.toml.tpl

@@ -0,0 +1,21 @@
+[build-system]
+requires = ["poetry-core>=1.4.0"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.poetry]
+name = "$project_name"
+version = "1.0.0"
+description = ""
+license = "Apache-2.0"
+authors = [
+    "The Flower Authors <[email protected]>",
+]
+readme = "README.md"
+
+[tool.poetry.dependencies]
+python = "^3.9"
+# Mandatory dependencies
+flwr-nightly = { version = "1.8.0.dev20240308", extras = ["simulation"] }
+flwr-datasets = { version = "^0.0.2", extras = ["vision"] }
+torch = "2.2.1"
+torchvision = "0.17.1"

src/py/flwr/cli/new/templates/app/requirements.numpy.txt.tpl

@@ -1,2 +1,2 @@
 flwr>=1.8, <2.0
-numpy >= 1.21.0
+numpy>=1.21.0

src/py/flwr/cli/new/templates/app/requirements.pytorch.txt.tpl

@@ -1,4 +1,4 @@
 flwr>=1.8, <2.0
 flwr-datasets[vision]>=0.0.2, <1.0.0
-torch==1.13.1
-torchvision==0.14.1
+torch==2.2.1
+torchvision==0.17.1

src/py/flwr/cli/new/templates/app/requirements.tensorflow.txt.tpl

@@ -1,4 +1,4 @@
 flwr>=1.8, <2.0
 flwr-datasets[vision]>=0.0.2, <1.0.0
-tensorflow-macos>=2.9.1, != 2.11.1 ; sys_platform == "darwin" and platform_machine == "arm64"
-tensorflow-cpu>=2.9.1, != 2.11.1 ; platform_machine == "x86_64"
+tensorflow-macos>=2.9.1, !=2.11.1 ; sys_platform == "darwin" and platform_machine == "arm64"
+tensorflow-cpu>=2.9.1, !=2.11.1 ; platform_machine == "x86_64"