Workspaces for different backends of Keras

.github/workflows/task_runner_basic_e2e.yml

@@ -36,7 +36,7 @@ jobs:
       matrix:
         # There are open issues for some of the models, so excluding them for now:
         # model_name: [ "torch_cnn_mnist", "keras_cnn_mnist", "torch_cnn_histology" ]
-        model_name: ["torch_cnn_mnist", "keras_cnn_mnist"]
+        model_name: ["torch_cnn_mnist", "keras/cnn_mnist"]
         python_version: ["3.10", "3.11", "3.12"]
       fail-fast: false # do not immediately fail if one of the combinations fail
 

.github/workflows/taskrunner.yml

@@ -42,4 +42,4 @@ jobs:
         pip install .
     - name: Test TaskRunner API
       run: |
-        python -m tests.github.test_hello_federation --template keras_cnn_mnist --fed_workspace aggregator --col1 col1 --col2 col2  --rounds-to-train 3 --save-model output_model
+        python -m tests.github.test_hello_federation --template keras/cnn_mnist --fed_workspace aggregator --col1 col1 --col2 col2  --rounds-to-train 3 --save-model output_model

.github/workflows/tr_docker_gramine_direct.yml

@@ -27,7 +27,7 @@ jobs:
 
     - name: Create workspace image
       run: |
-        fx workspace create --prefix example_workspace --template keras_cnn_mnist
+        fx workspace create --prefix example_workspace --template keras/cnn_mnist
         cd example_workspace
         fx plan initialize -a localhost
 

.github/workflows/tr_docker_native.yml

@@ -27,7 +27,7 @@ jobs:
 
     - name: Create workspace image
       run: |
-        fx workspace create --prefix example_workspace --template keras_cnn_mnist
+        fx workspace create --prefix example_workspace --template keras/cnn_mnist
         cd example_workspace
         fx plan initialize -a localhost
         fx workspace dockerize --save --revision https://github.com/${GITHUB_REPOSITORY}.git@${{ github.event.pull_request.head.sha }}

.github/workflows/ubuntu.yml

@@ -50,4 +50,4 @@ jobs:
           pip install .
       - name: Test TaskRunner API
         run: |
-          python -m tests.github.test_hello_federation --template keras_cnn_mnist --fed_workspace aggregator --col1 col1 --col2 col2  --rounds-to-train 3 --save-model output_model
+          python -m tests.github.test_hello_federation --template keras/cnn_mnist --fed_workspace aggregator --col1 col1 --col2 col2  --rounds-to-train 3 --save-model output_model

.github/workflows/windows.yml

@@ -48,4 +48,4 @@ jobs:
           pip install .
       - name: Test TaskRunner API
         run: |
-          python -m tests.github.test_hello_federation --template keras_cnn_mnist --fed_workspace aggregator --col1 col1 --col2 col2  --rounds-to-train 3 --save-model output_model
+          python -m tests.github.test_hello_federation --template keras/cnn_mnist --fed_workspace aggregator --col1 col1 --col2 col2  --rounds-to-train 3 --save-model output_model

Jenkinsfile

@@ -6,13 +6,12 @@ def snykData = [
     // CN-14619 snyk test CLI does not support -f in requirements.txt file
     // 'openfl-workspace_torch_cnn_histology': 'openfl-workspace/torch_cnn_histology/requirements.txt',
     'openfl-workspace_torch_cnn_histology_src': 'openfl-workspace/torch_cnn_histology/src/requirements.txt',
-    'openfl-workspace_keras_nlp': 'openfl-workspace/keras_nlp/requirements.txt',
+    'openfl-workspace_keras_nlp': 'openfl-workspace/keras/nlp/requirements.txt',
     'openfl-workspace_torch_cnn_mnist': 'openfl-workspace/torch_cnn_mnist/requirements.txt',
     'openfl-workspace_torch_unet_kvasir': 'openfl-workspace/torch_unet_kvasir/requirements.txt',
     'openfl-workspace_tf_cnn_histology': 'openfl-workspace/tf_cnn_histology/requirements.txt',
     'openfl-workspace_tf_3dunet_brats': 'openfl-workspace/tf_3dunet_brats/requirements.txt',
-    'openfl-workspace_keras_cnn_with_compression': 'openfl-workspace/keras_cnn_with_compression/requirements.txt',
-    'openfl-workspace_keras_cnn_mnist': 'openfl-workspace/keras_cnn_mnist/requirements.txt',
+    'openfl-workspace_keras_cnn_mnist': 'openfl-workspace/keras/cnn_mnist/requirements.txt',
     'openfl-tutorials_interactive_api_pytorch_medmnist_2d_envoy': 'openfl-tutorials/interactive_api/PyTorch_MedMNIST_2D/envoy/requirements.txt',
     'openfl-tutorials_interactive_api_pytorch_dogscats_vit_workspace': 'openfl-tutorials/interactive_api/PyTorch_DogsCats_ViT/workspace/requirements.txt',
     'openfl-tutorials_interactive_api_pytorch_histology_envoy': 'openfl-tutorials/interactive_api/PyTorch_Histology/envoy/requirements.txt',

docs/tutorials/taskrunner.ipynb

@@ -36,7 +36,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "!fx workspace create --prefix ./mnist_example --template keras_cnn_mnist\n",
+    "!fx workspace create --prefix ./mnist_example --template keras/cnn_mnist\n",
     "%cd ./mnist_example"
    ]
   },

openfl-workspace/keras/cnn_mnist/requirements.txt

@@ -0,0 +1,2 @@
+keras==3.6.0
+tensorflow==2.18.0

openfl-workspace/keras/nlp/.workspace

@@ -0,0 +1,2 @@
+current_plan_name: default
+

openfl-workspace/keras_jax/nlp/.workspace

@@ -0,0 +1,2 @@
+current_plan_name: default
+

openfl-workspace/keras_jax/nlp/plan/cols.yaml

@@ -0,0 +1,5 @@
+# Copyright (C) 2020-2021 Intel Corporation
+# Licensed subject to the terms of the separately executed evaluation license agreement between Intel Corporation and you.
+
+collaborators:
+

openfl-workspace/keras_jax/nlp/plan/data.yaml

@@ -0,0 +1,7 @@
+# Copyright (C) 2020-2021 Intel Corporation
+# Licensed subject to the terms of the separately executed evaluation license agreement between Intel Corporation and you.
+
+# collaborator_name,data_directory_path
+one,1
+
+

openfl-workspace/keras_jax/nlp/plan/plan.yaml

@@ -0,0 +1,46 @@
+# Copyright (C) 2020-2021 Intel Corporation
+# Licensed subject to the terms of the separately executed evaluation license agreement between Intel Corporation and you.
+
+aggregator :
+  defaults : plan/defaults/aggregator.yaml
+  template : openfl.component.Aggregator
+  settings :
+    init_state_path : save/init.pbuf
+    best_state_path : save/best.pbuf
+    last_state_path : save/last.pbuf
+    rounds_to_train : 10
+
+collaborator :
+  defaults : plan/defaults/collaborator.yaml
+  template : openfl.component.Collaborator
+  settings :
+    db_store_rounds: 2
+    delta_updates    : false
+    opt_treatment    : RESET
+
+data_loader :
+  defaults : plan/defaults/data_loader.yaml
+  template : src.dataloader.NLPDataLoader
+  settings :
+    collaborator_count : 2
+    batch_size         : 64
+    split_ratio: 0.2
+    num_samples: 10000
+
+task_runner :
+  defaults : plan/defaults/task_runner.yaml
+  template : src.taskrunner.KerasNLP
+  settings :
+    latent_dim : 256
+
+network :
+  defaults : plan/defaults/network.yaml
+
+assigner :
+  defaults : plan/defaults/assigner.yaml
+
+tasks :
+  defaults : plan/defaults/tasks_keras.yaml
+
+compression_pipeline :
+  defaults : plan/defaults/compression_pipeline.yaml

openfl-workspace/keras_jax/nlp/requirements.txt

@@ -0,0 +1,2 @@
+keras==3.6.0
+jax==0.4.38

openfl-workspace/keras_jax/nlp/src/__init__.py

@@ -0,0 +1,4 @@
+# Copyright (C) 2021-2022 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+"""openfl nlp keras template."""

openfl-workspace/keras_jax/nlp/src/dataloader.py

@@ -0,0 +1,142 @@
+"""Copyright (C) 2020-2021 Intel Corporation
+   SPDX-License-Identifier: Apache-2.0
+
+Licensed subject to the terms of the separately executed evaluation
+license agreement between Intel Corporation and you.
+"""
+from logging import getLogger
+from typing import Optional
+from typing import Iterator
+from typing import Tuple
+from typing import Union
+
+import numpy as np
+import src.dataloader_utils as dlu
+
+from openfl.federated import KerasDataLoader
+
+logger = getLogger(__name__)
+
+
+class NLPDataLoader(KerasDataLoader):
+    """NLP Dataloader template."""
+
+    def __init__(self, collaborator_count: int, split_ratio: float,
+                 num_samples: int, data_path: str, batch_size: int, **kwargs) -> None:
+        """Instantiate the data object.
+
+        Args:
+            data_path: The file path to the data Returns:
+            batch_size: The batch size of the data loader tuple: shape of an example feature array
+            **kwargs: Additional arguments, passed to super init and load_mnist_shard
+
+        Returns:
+           none
+        """
+        self.shard_num = data_path
+        self.data_path = dlu.download_data_()
+
+        self.batch_size = batch_size
+
+        train, valid, details = dlu.load_shard(collaborator_count, self.shard_num,
+                                               self.data_path, num_samples, split_ratio)
+
+        self.num_samples = details['num_samples']
+        self.num_encoder_tokens = details['num_encoder_tokens']
+        self.num_decoder_tokens = details['num_decoder_tokens']
+        self.max_encoder_seq_length = details['max_encoder_seq_length']
+        self.max_decoder_seq_length = details['max_decoder_seq_length']
+
+        self.X_train = [train[0], train[1]]
+        self.y_train = train[2]
+        self.X_valid = [valid[0], valid[1]]
+        self.y_valid = valid[2]
+
+    def get_feature_shape(self) -> Tuple[int, ...]:
+        """Get the shape of an example feature array."""
+        return self.X_train[0].shape
+
+    def get_train_loader(self, batch_size: Optional[int] = None) -> Iterator[Tuple[np.ndarray]]:
+        """
+        Get training data loader.
+
+        Returns
+        -------
+        loader object
+        """
+        return self._get_batch_generator(X1=self.X_train[0], X2=self.X_train[1],
+                                         y=self.y_train, batch_size=batch_size)
+
+    def get_valid_loader(self, batch_size: Optional[int] = None) -> Iterator[Tuple[np.ndarray]]:
+        """
+        Get validation data loader.
+
+        Returns:
+            loader object
+        """
+        return self._get_batch_generator(X1=self.X_valid[0], X2=self.X_valid[1],
+                                         y=self.y_valid, batch_size=batch_size)
+
+    def get_train_data_size(self) -> int:
+        """
+        Get total number of training samples.
+
+        Returns:
+            int: number of training samples
+        """
+        return self.X_train[0].shape[0]
+
+    def get_valid_data_size(self) -> int:
+        """
+        Get total number of validation samples.
+
+        Returns:
+            int: number of validation samples
+        """
+        return self.X_valid[0].shape[0]
+
+    @staticmethod
+    def _batch_generator(X1: np.ndarray, X2: np.ndarray,
+                         y: np.ndarray, idxs: np.ndarray,
+                         batch_size: int,
+                         num_batches: int) -> Iterator[Tuple[np.ndarray]]:
+        """
+        Generate batch of data.
+
+        Args:
+            X: input data
+            y: label data
+            idxs: The index of the dataset
+            batch_size: The batch size for the data loader
+            num_batches: The number of batches
+        Yields:
+            tuple: input data, label data
+        """
+        for i in range(num_batches):
+            a = i * batch_size
+            b = a + batch_size
+            yield (X1[idxs[a:b]], X2[idxs[a:b]]), y[idxs[a:b]]
+
+    def _get_batch_generator(self, X1: np.ndarray, X2: np.ndarray,
+                             y: np.ndarray,
+                             batch_size: Union[int, None]):
+        """
+        Return the dataset generator.
+
+        Args:
+            X1: input data  (encoder)
+            X2: input data  (decoder)
+            y: label data
+            batch_size: The batch size for the data loader
+        """
+        if batch_size is None:
+            batch_size = self.batch_size
+        # shuffle data indices
+        idxs = np.random.permutation(np.arange(X1.shape[0]))
+        # compute the number of batches
+        num_batches = int(np.ceil(X1.shape[0] / batch_size))
+        # build the generator and return it
+        # TODO: due to _batch_generator(X1, ...) has first param X1, all params here will be moved,
+        #       X1 -> X2, X2 -> y, y -> idxs, idxs -> batch_size, batch_size -> num_batches,
+        #       and num_batches -> should be unexpected in this function
+        return self._batch_generator(X1, X2, y, idxs, batch_size, num_batches)