Update TensorFlow Task Runner and related workspaces

openfl-workspace/tf_2dunet/src/tf_2dunet.py

@@ -5,12 +5,12 @@
 
 import tensorflow.compat.v1 as tf
 
-from openfl.federated import TensorFlowTaskRunner
+from openfl.federated import TensorFlowTaskRunner_v1
 
 tf.disable_v2_behavior()
 
 
-class TensorFlow2DUNet(TensorFlowTaskRunner):
+class TensorFlow2DUNet(TensorFlowTaskRunner_v1):
     """Initialize.
 
     Args:

openfl-workspace/tf_3dunet_brats/plan/plan.yaml

@@ -56,7 +56,7 @@ task_runner:
   template: src.tf_3dunet_model.TensorFlow3dUNet
 tasks:
   aggregated_model_validation:
-    function: validate
+    function: validate_task
     kwargs:
       apply: global
       batch_size: 4
@@ -65,7 +65,7 @@ tasks:
       - soft_dice_coef
   defaults: plan/defaults/tasks_tensorflow.yaml
   locally_tuned_model_validation:
-    function: validate
+    function: validate_task
     kwargs:
       apply: local
       batch_size: 4
@@ -74,7 +74,7 @@ tasks:
       - soft_dice_coef
   settings: {}
   train:
-    function: train
+    function: train_task
     kwargs:
       batch_size: 4
       epochs: 1

openfl-workspace/tf_3dunet_brats/src/tf_3dunet_model.py

@@ -5,14 +5,14 @@
 
 import tensorflow as tf
 
-from openfl.federated import KerasTaskRunner
+from openfl.federated import TensorFlowTaskRunner
 from .define_model import build_model
 from .define_model import dice_coef
 from .define_model import dice_loss
 from .define_model import soft_dice_coef
 
 
-class TensorFlow3dUNet(KerasTaskRunner):
+class TensorFlow3dUNet(TensorFlowTaskRunner):
     """Initialize.
 
     Args:

openfl-workspace/tf_cnn_histology/plan/plan.yaml

@@ -5,10 +5,10 @@ aggregator :
   defaults : plan/defaults/aggregator.yaml
   template : openfl.component.Aggregator
   settings :
-    init_state_path  : save/tf_cnn_histology_init.pbuf
-    last_state_path  : save/tf_cnn_histology_latest.pbuf
-    best_state_path  : save/tf_cnn_histology_best.pbuf
-    db_store_rounds: 2
+    init_state_path : save/tf_cnn_histology_init.pbuf
+    last_state_path : save/tf_cnn_histology_latest.pbuf
+    best_state_path : save/tf_cnn_histology_best.pbuf
+    db_store_rounds : 2
     rounds_to_train : 10
 
 collaborator :
@@ -21,7 +21,7 @@ collaborator :
 
 data_loader :
   defaults : plan/defaults/data_loader.yaml
-  template : src.tfhistology_inmemory.TensorFlowHistologyInMemory
+  template : src.dataloader.TensorFlowHistologyInMemory
   settings :
     batch_size: 64
     percent_train: 0.8
@@ -30,7 +30,7 @@ data_loader :
 
 task_runner :
   defaults : plan/defaults/task_runner.yaml
-  template : src.tf_cnn.TensorFlowCNN
+  template : src.taskrunner.TensorFlowCNN
 
 network :
   defaults : plan/defaults/network.yaml
@@ -41,22 +41,22 @@ assigner :
 tasks:
   defaults: plan/defaults/tasks_tensorflow.yaml
   aggregated_model_validation:
-    function: validate
+    function: validate_task
     kwargs:
       apply: global
       batch_size: 32
       metrics:
       - sparse_categorical_accuracy
   locally_tuned_model_validation:
-    function: validate
+    function: validate_task
     kwargs:
       apply: local
       batch_size: 32
       metrics:
       - sparse_categorical_accuracy
   settings: {}
   train:
-    function: train
+    function: train_task
     kwargs:
       batch_size: 32
       epochs: 1

openfl-workspace/tf_cnn_histology/requirements.txt

@@ -1,3 +1,3 @@
 pillow
-tensorflow==2.13
+tensorflow==2.15.1
 tensorflow-datasets

openfl-workspace/tf_cnn_histology/src/tfds_utils.py → openfl-workspace/tf_cnn_histology/src/dataloader.py

@@ -3,13 +3,40 @@
 
 """You may copy this file as the starting point of your own model."""
 
+from openfl.federated import TensorFlowDataLoader
 from logging import getLogger
 
 import numpy as np
 import tensorflow_datasets as tfds
 
 logger = getLogger(__name__)
 
+class TensorFlowHistologyInMemory(TensorFlowDataLoader):
+    """TensorFlow Data Loader for Colorectal Histology Dataset."""
+
+    def __init__(self, data_path, batch_size, **kwargs):
+        """
+        Initialize.
+
+        Args:
+            data_path: File path for the dataset
+            batch_size (int): The batch size for the data loader
+            **kwargs: Additional arguments, passed to super init and load_mnist_shard
+        """
+        super().__init__(batch_size, **kwargs)
+
+        _, num_classes, X_train, y_train, X_valid, y_valid = load_histology_shard(
+            shard_num=data_path,
+            categorical=False, **kwargs
+        )
+
+        self.X_train = X_train
+        self.y_train = y_train
+        self.X_valid = X_valid
+        self.y_valid = y_valid
+
+        self.num_classes = num_classes
+
 
 def one_hot(labels, classes):
     """

openfl-workspace/tf_cnn_histology/src/tf_cnn.py → openfl-workspace/tf_cnn_histology/src/taskrunner.py

@@ -3,12 +3,14 @@
 
 """You may copy this file as the starting point of your own model."""
 
+import numpy as np
 import tensorflow as tf
 
-from openfl.federated import KerasTaskRunner
+from openfl.utilities import Metric
+from openfl.federated import TensorFlowTaskRunner
 
 
-class TensorFlowCNN(KerasTaskRunner):
+class TensorFlowCNN(TensorFlowTaskRunner):
     """Initialize.
 
     Args:
@@ -17,40 +19,47 @@ class TensorFlowCNN(KerasTaskRunner):
     """
 
     def __init__(self, **kwargs):
-        """Initialize.
+        """
+        Initialize.
 
         Args:
             **kwargs: Additional parameters to pass to the function
 
         """
         super().__init__(**kwargs)
 
-        self.model = self.create_model(
+        self.model = self.build_model(
             self.feature_shape,
             self.data_loader.num_classes,
             **kwargs
         )
         self.initialize_tensorkeys_for_functions()
 
-    def create_model(self,
-                     input_shape,
-                     num_classes,
-                     training_smoothing=32.0,
-                     validation_smoothing=1.0,
-                     **kwargs):
-        """Create the TensorFlow CNN Histology model.
+        self.model.summary(print_fn=self.logger.info)
+
+        self.logger.info(f'Train Set Size : {self.get_train_data_size()}')
+        self.logger.info(f'Valid Set Size : {self.get_valid_data_size()}')
+
+    def build_model(self,
+                    input_shape,
+                    num_classes,
+                    **kwargs):
+        """
+        Define the model architecture.
 
         Args:
-            training_smoothing (float): (Default=32.0)
-            validation_smoothing (float): (Default=1.0)
+            input_shape (numpy.ndarray): The shape of the data
+            num_classes (int): The number of classes of the dataset
             **kwargs: Additional parameters to pass to the function
 
+        Returns:
+            keras.src.engine.functional.Functional
+
         """
         print(tf.config.threading.get_intra_op_parallelism_threads())
         print(tf.config.threading.get_inter_op_parallelism_threads())
-        # ## Define Model
-        #
-        # Convolutional neural network model
+
+        ## Define Model using Functional API
 
         inputs = tf.keras.layers.Input(shape=input_shape)
         conv = tf.keras.layers.Conv2D(
@@ -96,13 +105,39 @@ def create_model(self,
             metrics=[tf.keras.metrics.SparseCategoricalAccuracy()],
         )
 
-        self.tvars = model.layers
-        print(f'layer names: {[var.name for var in self.tvars]}')
+        return model
 
-        self.opt_vars = self.optimizer.variables()
-        print(f'optimizer vars: {self.opt_vars}')
 
-        # Two opt_vars for one tvar: gradient and square sum for RMSprop.
-        self.fl_vars = self.tvars + self.opt_vars
+    def train_(self, batch_generator, metrics: list = None, **kwargs):
+        """Train single epoch.
 
-        return model
+        Override this function for custom training.
+
+        Args:
+            batch_generator: Generator of training batches.
+                Each batch is a tuple of N train images and N train labels
+                where N is the batch size of the DataLoader of the current TaskRunner instance.
+
+            epochs: Number of epochs to train.
+            metrics: Names of metrics to save.
+        """
+        if metrics is None:
+            metrics = []
+
+        model_metrics_names = self.model.metrics_names
+
+        for param in metrics:
+            if param not in model_metrics_names:
+                raise ValueError(
+                    f'TensorFlowTaskRunner does not support specifying new metrics. '
+                    f'Param_metrics = {metrics}, model_metrics_names = {model_metrics_names}'
+                )
+
+        history = self.model.fit(batch_generator,
+                                 verbose=1,
+                                 **kwargs)
+        results = []
+        for metric in metrics:
+            value = np.mean([history.history[metric]])
+            results.append(Metric(name=metric, value=np.array(value)))
+        return results

openfl-workspace/tf_cnn_mnist/.workspace

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

openfl-workspace/tf_cnn_mnist/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/tf_cnn_mnist/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/tf_cnn_mnist/plan/defaults

@@ -0,0 +1,2 @@
+../../workspace/plan/defaults
+

openfl-workspace/tf_cnn_mnist/plan/plan.yaml

@@ -0,0 +1,42 @@
+# Copyright (C) 2020-2024 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/tf_cnn_mnist_init.pbuf
+    best_state_path : save/tf_cnn_mnist_best.pbuf
+    last_state_path : save/tf_cnn_mnist_last.pbuf
+    rounds_to_train : 10
+
+collaborator :
+  defaults : plan/defaults/collaborator.yaml
+  template : openfl.component.Collaborator
+  settings :
+    delta_updates    : false
+    opt_treatment    : RESET
+
+data_loader :
+  defaults : plan/defaults/data_loader.yaml
+  template : src.dataloader.TensorFlowMNISTInMemory
+  settings :
+    collaborator_count : 2
+    data_group_name    : mnist
+    batch_size         : 256
+
+task_runner :
+  defaults : plan/defaults/task_runner.yaml
+  template : src.taskrunner.TensorFlowCNN
+
+network :
+  defaults : plan/defaults/network.yaml
+
+assigner :
+  defaults : plan/defaults/assigner.yaml
+
+tasks :
+  defaults : plan/defaults/tasks_tensorflow.yaml
+
+compression_pipeline :
+  defaults : plan/defaults/compression_pipeline.yaml

openfl-workspace/tf_cnn_mnist/requirements.txt

@@ -0,0 +1 @@
+tensorflow==2.15.1

openfl-workspace/tf_cnn_mnist/src/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (C) 2020-2021 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+"""You may copy this file as the starting point of your own model."""