flwr-nightly 1.11.0.dev20240724__tar.gz → 1.11.0.dev20240811__tar.gz

{flwr_nightly-1.11.0.dev20240724 → flwr_nightly-1.11.0.dev20240811}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flwr-nightly
-Version: 1.11.0.dev20240724
+Version: 1.11.0.dev20240811
 Summary: Flower: A Friendly Federated Learning Framework
 Home-page: https://flower.ai
 License: Apache-2.0
@@ -33,7 +33,7 @@ Classifier: Typing :: Typed
 Provides-Extra: rest
 Provides-Extra: simulation
 Requires-Dist: cryptography (>=42.0.4,<43.0.0)
-Requires-Dist: grpcio (>=1.60.0,<2.0.0,!=1.64.2,!=1.65.1)
+Requires-Dist: grpcio (>=1.60.0,<2.0.0,!=1.64.2,!=1.65.1,!=1.65.2,!=1.65.4)
 Requires-Dist: iterators (>=0.0.2,<0.0.3)
 Requires-Dist: numpy (>=1.21.0,<2.0.0)
 Requires-Dist: pathspec (>=0.12.1,<0.13.0)

{flwr_nightly-1.11.0.dev20240724 → flwr_nightly-1.11.0.dev20240811}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "flwr-nightly"
-version = "1.11.0.dev20240724"
+version = "1.11.0.dev20240811"
 description = "Flower: A Friendly Federated Learning Framework"
 license = "Apache-2.0"
 authors = ["The Flower Authors <hello@flower.ai>"]
@@ -64,12 +64,12 @@ flower-simulation = "flwr.simulation.run_simulation:run_simulation_from_cli"
 python = "^3.8"
 # Mandatory dependencies
 numpy = "^1.21.0"
-grpcio = "^1.60.0,!=1.64.2,!=1.65.1"
+grpcio = "^1.60.0,!=1.64.2,!=1.65.1,!=1.65.2,!=1.65.4"
 protobuf = "^4.25.2"
 cryptography = "^42.0.4"
 pycryptodome = "^3.18.0"
 iterators = "^0.0.2"
-typer = { version = "^0.9.0", extras=["all"] }
+typer = { version = "^0.9.0", extras = ["all"] }
 tomli = "^2.0.1"
 tomli-w = "^1.0.0"
 pathspec = "^0.12.1"

{flwr_nightly-1.11.0.dev20240724 → flwr_nightly-1.11.0.dev20240811}/src/py/flwr/cli/build.py

@@ -30,32 +30,34 @@ from .utils import get_sha256_hash, is_valid_project_name
 
 # pylint: disable=too-many-locals
 def build(
-    directory: Annotated[
+    app: Annotated[
         Optional[Path],
-        typer.Option(help="Path of the Flower project to bundle into a FAB"),
+        typer.Option(help="Path of the Flower App to bundle into a FAB"),
     ] = None,
 ) -> str:
-    """Build a Flower project into a Flower App Bundle (FAB).
+    """Build a Flower App into a Flower App Bundle (FAB).
 
-    You can run ``flwr build`` without any arguments to bundle the current directory,
-    or you can use ``--directory`` to build a specific directory:
-    ``flwr build --directory ./projects/flower-hello-world``.
+    You can run ``flwr build`` without any arguments to bundle the app located in the
+    current directory. Alternatively, you can you can specify a path using the ``--app``
+    option to bundle an app located at the provided path. For example:
+
+    ``flwr build --app ./apps/flower-hello-world``.
     """
-    if directory is None:
-        directory = Path.cwd()
+    if app is None:
+        app = Path.cwd()
 
-    directory = directory.resolve()
-    if not directory.is_dir():
+    app = app.resolve()
+    if not app.is_dir():
         typer.secho(
-            f"❌ The path {directory} is not a valid directory.",
+            f"❌ The path {app} is not a valid path to a Flower app.",
             fg=typer.colors.RED,
             bold=True,
         )
         raise typer.Exit(code=1)
 
-    if not is_valid_project_name(directory.name):
+    if not is_valid_project_name(app.name):
         typer.secho(
-            f"❌ The project name {directory.name} is invalid, "
+            f"❌ The project name {app.name} is invalid, "
             "a valid project name must start with a letter or an underscore, "
             "and can only contain letters, digits, and underscores.",
             fg=typer.colors.RED,
@@ -63,7 +65,7 @@ def build(
         )
         raise typer.Exit(code=1)
 
-    conf, errors, warnings = load_and_validate(directory / "pyproject.toml")
+    conf, errors, warnings = load_and_validate(app / "pyproject.toml")
     if conf is None:
         typer.secho(
             "Project configuration could not be loaded.\npyproject.toml is invalid:\n"
@@ -82,12 +84,12 @@ def build(
         )
 
     # Load .gitignore rules if present
-    ignore_spec = _load_gitignore(directory)
+    ignore_spec = _load_gitignore(app)
 
     # Set the name of the zip file
     fab_filename = (
         f"{conf['tool']['flwr']['app']['publisher']}"
-        f".{directory.name}"
+        f".{app.name}"
         f".{conf['project']['version'].replace('.', '-')}.fab"
     )
     list_file_content = ""
@@ -108,7 +110,7 @@ def build(
         fab_file.writestr("pyproject.toml", toml_contents)
 
         # Continue with adding other files
-        for root, _, files in os.walk(directory, topdown=True):
+        for root, _, files in os.walk(app, topdown=True):
             files = [
                 f
                 for f in files
@@ -120,7 +122,7 @@ def build(
 
             for file in files:
                 file_path = Path(root) / file
-                archive_path = file_path.relative_to(directory)
+                archive_path = file_path.relative_to(app)
                 fab_file.write(file_path, archive_path)
 
                 # Calculate file info
@@ -138,9 +140,9 @@ def build(
     return fab_filename
 
 
-def _load_gitignore(directory: Path) -> pathspec.PathSpec:
+def _load_gitignore(app: Path) -> pathspec.PathSpec:
     """Load and parse .gitignore file, returning a pathspec."""
-    gitignore_path = directory / ".gitignore"
+    gitignore_path = app / ".gitignore"
     patterns = ["__pycache__/"]  # Default pattern
     if gitignore_path.exists():
         with open(gitignore_path, encoding="UTF-8") as file:

{flwr_nightly-1.11.0.dev20240724 → flwr_nightly-1.11.0.dev20240811}/src/py/flwr/cli/config_utils.py

@@ -25,8 +25,8 @@ from flwr.common import object_ref
 from flwr.common.typing import UserConfigValue
 
 
-def get_fab_metadata(fab_file: Union[Path, bytes]) -> Tuple[str, str]:
-    """Extract the fab_id and the fab_version from a FAB file or path.
+def get_fab_config(fab_file: Union[Path, bytes]) -> Dict[str, Any]:
+    """Extract the config from a FAB file or path.
 
     Parameters
     ----------
@@ -36,8 +36,8 @@ def get_fab_metadata(fab_file: Union[Path, bytes]) -> Tuple[str, str]:
 
     Returns
     -------
-    Tuple[str, str]
-        The `fab_version` and `fab_id` of the given Flower App Bundle.
+    Dict[str, Any]
+        The `config` of the given Flower App Bundle.
     """
     fab_file_archive: Union[Path, IO[bytes]]
     if isinstance(fab_file, bytes):
@@ -59,10 +59,29 @@ def get_fab_metadata(fab_file: Union[Path, bytes]) -> Tuple[str, str]:
         if not is_valid:
             raise ValueError(errors)
 
-        return (
-            conf["project"]["version"],
-            f"{conf['tool']['flwr']['app']['publisher']}/{conf['project']['name']}",
-        )
+        return conf
+
+
+def get_fab_metadata(fab_file: Union[Path, bytes]) -> Tuple[str, str]:
+    """Extract the fab_id and the fab_version from a FAB file or path.
+
+    Parameters
+    ----------
+    fab_file : Union[Path, bytes]
+        The Flower App Bundle file to validate and extract the metadata from.
+        It can either be a path to the file or the file itself as bytes.
+
+    Returns
+    -------
+    Tuple[str, str]
+        The `fab_version` and `fab_id` of the given Flower App Bundle.
+    """
+    conf = get_fab_config(fab_file)
+
+    return (
+        conf["project"]["version"],
+        f"{conf['tool']['flwr']['app']['publisher']}/{conf['project']['name']}",
+    )
 
 
 def load_and_validate(

{flwr_nightly-1.11.0.dev20240724 → flwr_nightly-1.11.0.dev20240811}/src/py/flwr/cli/new/new.py

@@ -34,13 +34,13 @@ from ..utils import (
 class MlFramework(str, Enum):
     """Available frameworks."""
 
-    NUMPY = "NumPy"
     PYTORCH = "PyTorch"
     TENSORFLOW = "TensorFlow"
-    JAX = "JAX"
+    SKLEARN = "sklearn"
     HUGGINGFACE = "HuggingFace"
+    JAX = "JAX"
     MLX = "MLX"
-    SKLEARN = "sklearn"
+    NUMPY = "NumPy"
     FLOWERTUNE = "FlowerTune"
 
 
@@ -92,9 +92,9 @@ def render_and_create(file_path: Path, template: str, context: Dict[str, str]) -
 
 # pylint: disable=too-many-locals,too-many-branches,too-many-statements
 def new(
-    project_name: Annotated[
+    app_name: Annotated[
         Optional[str],
-        typer.Argument(metavar="project_name", help="The name of the project"),
+        typer.Argument(help="The name of the Flower App"),
     ] = None,
     framework: Annotated[
         Optional[MlFramework],
@@ -105,26 +105,26 @@ def new(
         typer.Option(case_sensitive=False, help="The Flower username of the author"),
     ] = None,
 ) -> None:
-    """Create new Flower project."""
-    if project_name is None:
-        project_name = prompt_text("Please provide the project name")
-    if not is_valid_project_name(project_name):
-        project_name = prompt_text(
+    """Create new Flower App."""
+    if app_name is None:
+        app_name = prompt_text("Please provide the app name")
+    if not is_valid_project_name(app_name):
+        app_name = prompt_text(
             "Please provide a name that only contains "
             "characters in {'-', a-zA-Z', '0-9'}",
             predicate=is_valid_project_name,
-            default=sanitize_project_name(project_name),
+            default=sanitize_project_name(app_name),
         )
 
     # Set project directory path
-    package_name = re.sub(r"[-_.]+", "-", project_name).lower()
+    package_name = re.sub(r"[-_.]+", "-", app_name).lower()
     import_name = package_name.replace("-", "_")
     project_dir = Path.cwd() / package_name
 
     if project_dir.exists():
         if not typer.confirm(
             typer.style(
-                f"\n💬 {project_name} already exists, do you want to override it?",
+                f"\n💬 {app_name} already exists, do you want to override it?",
                 fg=typer.colors.MAGENTA,
                 bold=True,
             )
@@ -135,20 +135,20 @@ def new(
         username = prompt_text("Please provide your Flower username")
 
     if framework is not None:
-        framework_str = str(framework.value)
+        framework_str_upper = str(framework.value)
     else:
         framework_value = prompt_options(
             "Please select ML framework by typing in the number",
-            sorted([mlf.value for mlf in MlFramework]),
+            [mlf.value for mlf in MlFramework],
         )
         selected_value = [
             name
             for name, value in vars(MlFramework).items()
             if value == framework_value
         ]
-        framework_str = selected_value[0]
+        framework_str_upper = selected_value[0]
 
-    framework_str = framework_str.lower()
+    framework_str = framework_str_upper.lower()
 
     llm_challenge_str = None
     if framework_str == "flowertune":
@@ -166,16 +166,17 @@ def new(
 
     print(
         typer.style(
-            f"\n🔨 Creating Flower project {project_name}...",
+            f"\n🔨 Creating Flower App {app_name}...",
             fg=typer.colors.GREEN,
             bold=True,
         )
     )
 
     context = {
-        "project_name": project_name,
-        "package_name": package_name,
+        "framework_str": framework_str_upper,
         "import_name": import_name.replace("-", "_"),
+        "package_name": package_name,
+        "project_name": app_name,
         "username": username,
     }
 
@@ -267,8 +268,8 @@ def new(
 
     print(
         typer.style(
-            "🎊 Project creation successful.\n\n"
-            "Use the following command to run your project:\n",
+            "🎊 Flower App creation successful.\n\n"
+            "Use the following command to run your Flower App:\n",
             fg=typer.colors.GREEN,
             bold=True,
         )

{flwr_nightly-1.11.0.dev20240724 → flwr_nightly-1.11.0.dev20240811}/src/py/flwr/cli/new/templates/app/README.md.tpl

@@ -1,4 +1,4 @@
-# $project_name
+# $project_name: A Flower / $framework_str app
 
 ## Install dependencies
 

flwr_nightly-1.11.0.dev20240811/src/py/flwr/cli/new/templates/app/code/__init__.py.tpl

@@ -0,0 +1 @@
+"""$project_name: A Flower / $framework_str app."""

{flwr_nightly-1.11.0.dev20240724 → flwr_nightly-1.11.0.dev20240811}/src/py/flwr/cli/new/templates/app/code/client.huggingface.py.tpl

@@ -1,4 +1,4 @@
-"""$project_name: A Flower / HuggingFace Transformers app."""
+"""$project_name: A Flower / $framework_str app."""
 
 from flwr.client import ClientApp, NumPyClient
 from flwr.common import Context

{flwr_nightly-1.11.0.dev20240724 → flwr_nightly-1.11.0.dev20240811}/src/py/flwr/cli/new/templates/app/code/client.jax.py.tpl

@@ -1,4 +1,4 @@
-"""$project_name: A Flower / JAX app."""
+"""$project_name: A Flower / $framework_str app."""
 
 import jax
 from flwr.client import NumPyClient, ClientApp

{flwr_nightly-1.11.0.dev20240724 → flwr_nightly-1.11.0.dev20240811}/src/py/flwr/cli/new/templates/app/code/client.mlx.py.tpl

@@ -1,4 +1,4 @@
-"""$project_name: A Flower / MLX app."""
+"""$project_name: A Flower / $framework_str app."""
 
 import mlx.core as mx
 import mlx.nn as nn

{flwr_nightly-1.11.0.dev20240724 → flwr_nightly-1.11.0.dev20240811}/src/py/flwr/cli/new/templates/app/code/client.numpy.py.tpl

@@ -1,4 +1,4 @@
-"""$project_name: A Flower / NumPy app."""
+"""$project_name: A Flower / $framework_str app."""
 
 from flwr.client import NumPyClient, ClientApp
 from flwr.common import Context

{flwr_nightly-1.11.0.dev20240724 → flwr_nightly-1.11.0.dev20240811}/src/py/flwr/cli/new/templates/app/code/client.pytorch.py.tpl

@@ -1,11 +1,11 @@
-"""$project_name: A Flower / PyTorch app."""
+"""$project_name: A Flower / $framework_str app."""
 
+import torch
 from flwr.client import NumPyClient, ClientApp
 from flwr.common import Context
 
 from $import_name.task import (
     Net,
-    DEVICE,
     load_data,
     get_weights,
     set_weights,
@@ -21,27 +21,28 @@ class FlowerClient(NumPyClient):
         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)
-        results = train(
+        train_loss = train(
             self.net,
             self.trainloader,
-            self.valloader,
             self.local_epochs,
-            DEVICE,
+            self.device,
         )
-        return get_weights(self.net), len(self.trainloader.dataset), results
+        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)
+        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().to(DEVICE)
+    net = Net()
     partition_id = context.node_config["partition-id"]
     num_partitions = context.node_config["num-partitions"]
     trainloader, valloader = load_data(partition_id, num_partitions)

{flwr_nightly-1.11.0.dev20240724 → flwr_nightly-1.11.0.dev20240811}/src/py/flwr/cli/new/templates/app/code/client.sklearn.py.tpl

@@ -1,4 +1,4 @@
-"""$project_name: A Flower / Scikit-Learn app."""
+"""$project_name: A Flower / $framework_str app."""
 
 import warnings
 

{flwr_nightly-1.11.0.dev20240724 → flwr_nightly-1.11.0.dev20240811}/src/py/flwr/cli/new/templates/app/code/client.tensorflow.py.tpl

@@ -1,4 +1,4 @@
-"""$project_name: A Flower / TensorFlow app."""
+"""$project_name: A Flower / $framework_str app."""
 
 from flwr.client import NumPyClient, ClientApp
 from flwr.common import Context
@@ -9,13 +9,10 @@ from $import_name.task import load_data, load_model
 # Define Flower Client and client_fn
 class FlowerClient(NumPyClient):
     def __init__(
-        self, model, x_train, y_train, x_test, y_test, epochs, batch_size, verbose
+        self, model, data, epochs, batch_size, verbose
     ):
         self.model = model
-        self.x_train = x_train
-        self.y_train = y_train
-        self.x_test = x_test
-        self.y_test = y_test
+        self.x_train, self.y_train, self.x_test, self.y_test = data
         self.epochs = epochs
         self.batch_size = batch_size
         self.verbose = verbose
@@ -46,14 +43,14 @@ def client_fn(context: Context):
 
     partition_id = context.node_config["partition-id"]
     num_partitions = context.node_config["num-partitions"]
-    x_train, y_train, x_test, y_test = load_data(partition_id, num_partitions)
+    data = load_data(partition_id, num_partitions)
     epochs = context.run_config["local-epochs"]
     batch_size = context.run_config["batch-size"]
     verbose = context.run_config.get("verbose")
 
     # Return Client instance
     return FlowerClient(
-        net, x_train, y_train, x_test, y_test, epochs, batch_size, verbose
+        net, data, epochs, batch_size, verbose
     ).to_client()
 
 

{flwr_nightly-1.11.0.dev20240724 → flwr_nightly-1.11.0.dev20240811}/src/py/flwr/cli/new/templates/app/code/server.huggingface.py.tpl

@@ -1,4 +1,4 @@
-"""$project_name: A Flower / HuggingFace Transformers app."""
+"""$project_name: A Flower / $framework_str app."""
 
 from flwr.common import Context
 from flwr.server.strategy import FedAvg

{flwr_nightly-1.11.0.dev20240724 → flwr_nightly-1.11.0.dev20240811}/src/py/flwr/cli/new/templates/app/code/server.jax.py.tpl

@@ -1,4 +1,4 @@
-"""$project_name: A Flower / JAX app."""
+"""$project_name: A Flower / $framework_str app."""
 
 from flwr.common import Context
 from flwr.server.strategy import FedAvg

{flwr_nightly-1.11.0.dev20240724 → flwr_nightly-1.11.0.dev20240811}/src/py/flwr/cli/new/templates/app/code/server.mlx.py.tpl

@@ -1,4 +1,4 @@
-"""$project_name: A Flower / MLX app."""
+"""$project_name: A Flower / $framework_str app."""
 
 from flwr.common import Context
 from flwr.server import ServerApp, ServerAppComponents, ServerConfig

{flwr_nightly-1.11.0.dev20240724 → flwr_nightly-1.11.0.dev20240811}/src/py/flwr/cli/new/templates/app/code/server.numpy.py.tpl

@@ -1,4 +1,4 @@
-"""$project_name: A Flower / NumPy app."""
+"""$project_name: A Flower / $framework_str app."""
 
 from flwr.common import Context
 from flwr.server import ServerApp, ServerAppComponents, ServerConfig

{flwr_nightly-1.11.0.dev20240724 → flwr_nightly-1.11.0.dev20240811}/src/py/flwr/cli/new/templates/app/code/server.pytorch.py.tpl

@@ -1,4 +1,4 @@
-"""$project_name: A Flower / PyTorch app."""
+"""$project_name: A Flower / $framework_str app."""
 
 from flwr.common import Context, ndarrays_to_parameters
 from flwr.server import ServerApp, ServerAppComponents, ServerConfig
@@ -7,17 +7,18 @@ from flwr.server.strategy import FedAvg
 from $import_name.task import Net, get_weights
 
 
-# Initialize model parameters
-ndarrays = get_weights(Net())
-parameters = ndarrays_to_parameters(ndarrays)
-
 def server_fn(context: Context):
     # Read from config
     num_rounds = context.run_config["num-server-rounds"]
+    fraction_fit = context.run_config["fraction-fit"]
+
+    # Initialize model parameters
+    ndarrays = get_weights(Net())
+    parameters = ndarrays_to_parameters(ndarrays)
 
     # Define strategy
     strategy = FedAvg(
-        fraction_fit=1.0,
+        fraction_fit=fraction_fit,
         fraction_evaluate=1.0,
         min_available_clients=2,
         initial_parameters=parameters,

{flwr_nightly-1.11.0.dev20240724 → flwr_nightly-1.11.0.dev20240811}/src/py/flwr/cli/new/templates/app/code/server.sklearn.py.tpl

@@ -1,4 +1,4 @@
-"""$project_name: A Flower / Scikit-Learn app."""
+"""$project_name: A Flower / $framework_str app."""
 
 from flwr.common import Context
 from flwr.server import ServerApp, ServerAppComponents, ServerConfig

{flwr_nightly-1.11.0.dev20240724 → flwr_nightly-1.11.0.dev20240811}/src/py/flwr/cli/new/templates/app/code/server.tensorflow.py.tpl

@@ -1,4 +1,4 @@
-"""$project_name: A Flower / TensorFlow app."""
+"""$project_name: A Flower / $framework_str app."""
 
 from flwr.common import Context, ndarrays_to_parameters
 from flwr.server import ServerApp, ServerAppComponents, ServerConfig
@@ -6,15 +6,14 @@ from flwr.server.strategy import FedAvg
 
 from $import_name.task import load_model
 
-# Define config
-config = ServerConfig(num_rounds=3)
-
-parameters = ndarrays_to_parameters(load_model().get_weights())
 
 def server_fn(context: Context):
     # Read from config
     num_rounds = context.run_config["num-server-rounds"]
 
+    # Get parameters to initialize global model
+    parameters = ndarrays_to_parameters(load_model().get_weights())
+
     # Define strategy
     strategy = strategy = FedAvg(
         fraction_fit=1.0,

{flwr_nightly-1.11.0.dev20240724 → flwr_nightly-1.11.0.dev20240811}/src/py/flwr/cli/new/templates/app/code/task.huggingface.py.tpl

@@ -1,4 +1,4 @@
-"""$project_name: A Flower / HuggingFace Transformers app."""
+"""$project_name: A Flower / $framework_str app."""
 
 import warnings
 from collections import OrderedDict

{flwr_nightly-1.11.0.dev20240724 → flwr_nightly-1.11.0.dev20240811}/src/py/flwr/cli/new/templates/app/code/task.jax.py.tpl

@@ -1,4 +1,4 @@
-"""$project_name: A Flower / JAX app."""
+"""$project_name: A Flower / $framework_str app."""
 
 import jax
 import jax.numpy as jnp
@@ -33,7 +33,7 @@ def train(params, grad_fn, X, y):
     num_examples = X.shape[0]
     for epochs in range(50):
         grads = grad_fn(params, X, y)
-        params = jax.tree.map(lambda p, g: p - 0.05 * g, params, grads)
+        params = jax.tree_map(lambda p, g: p - 0.05 * g, params, grads)
         loss = loss_fn(params, X, y)
     return params, loss, num_examples
 

{flwr_nightly-1.11.0.dev20240724 → flwr_nightly-1.11.0.dev20240811}/src/py/flwr/cli/new/templates/app/code/task.mlx.py.tpl

@@ -1,4 +1,4 @@
-"""$project_name: A Flower / MLX app."""
+"""$project_name: A Flower / $framework_str app."""
 
 import mlx.core as mx
 import mlx.nn as nn
@@ -56,6 +56,7 @@ def load_data(partition_id: int, num_partitions: int):
         fds = FederatedDataset(
             dataset="ylecun/mnist",
             partitioners={"train": partitioner},
+            trust_remote_code=True,
         )
     partition = fds.load_partition(partition_id)
     partition_splits = partition.train_test_split(test_size=0.2, seed=42)

{flwr_nightly-1.11.0.dev20240724 → flwr_nightly-1.11.0.dev20240811}/src/py/flwr/cli/new/templates/app/code/task.pytorch.py.tpl

@@ -1,4 +1,4 @@
-"""$project_name: A Flower / PyTorch app."""
+"""$project_name: A Flower / $framework_str app."""
 
 from collections import OrderedDict
 
@@ -11,9 +11,6 @@ from flwr_datasets import FederatedDataset
 from flwr_datasets.partitioner import IidPartitioner
 
 
-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')"""
 
@@ -66,44 +63,41 @@ def load_data(partition_id: int, num_partitions: int):
     return trainloader, testloader
 
 
-def train(net, trainloader, valloader, epochs, device):
+def train(net, trainloader, epochs, device):
     """Train the model on the training set."""
     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)
+    optimizer = torch.optim.SGD(net.parameters(), lr=0.1, momentum=0.9)
     net.train()
+    running_loss = 0.0
     for _ in range(epochs):
         for batch in trainloader:
             images = batch["img"]
             labels = batch["label"]
             optimizer.zero_grad()
-            criterion(net(images.to(DEVICE)), labels.to(DEVICE)).backward()
+            loss = criterion(net(images.to(device)), labels.to(device))
+            loss.backward()
             optimizer.step()
+            running_loss += loss.item()
 
-    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
+    avg_trainloss = running_loss / len(trainloader)
+    return avg_trainloss
 
 
-def test(net, testloader):
+def test(net, testloader, device):
     """Validate the model on the test set."""
+    net.to(device)
     criterion = torch.nn.CrossEntropyLoss()
     correct, loss = 0, 0.0
     with torch.no_grad():
         for batch in testloader:
-            images = batch["img"].to(DEVICE)
-            labels = batch["label"].to(DEVICE)
+            images = batch["img"].to(device)
+            labels = batch["label"].to(device)
             outputs = net(images)
             loss += criterion(outputs, labels).item()
             correct += (torch.max(outputs.data, 1)[1] == labels).sum().item()
     accuracy = correct / len(testloader.dataset)
+    loss = loss / len(testloader)
     return loss, accuracy