Updated Keras and TensorFlow Task Runner and related workspaces.

docs/developer_ref/troubleshooting.rst

@@ -10,7 +10,7 @@
 The following is a list of commonly reported issues in Open Federated Learning (|productName|). If you don't see your issue reported here, please submit a `Github issue
 <https://github.com/intel/openfl/issues>`_ or contact us directly on `Slack <https://join.slack.com/t/openfl/shared_invite/zt-ovzbohvn-T5fApk05~YS_iZhjJ5yaTw>`_.
 
-1. I see the error :code:`Cannot import name TensorFlowDataLoader from openfl.federated`
+1. I see the error :code:`Cannot import name KerasDataLoader from openfl.federated`
 
    |productName| currently uses conditional imports to attempt to be framework agnostic. If your task runner is derived from `KerasTaskRunner` or `TensorflowTaskRunner`, this error could come up if TensorFlow\*\  was not installed in your collaborator's virtual environment. If running on multi-node experiment, we recommend using the :code:`fx workspace export` and :code:`fx workspace import` commands, as this will ensure consistent modules between aggregator and collaborators.
 

openfl-workspace/tf_cnn_histology/plan/data.yaml → openfl-workspace/keras_2dunet/plan/data.yaml

@@ -1,4 +1,4 @@
-# Copyright (C) 2020-2021 Intel Corporation
+# Copyright (C) 2020-2024 Intel Corporation
 # Licensed subject to the terms of the separately executed evaluation license agreement between Intel Corporation and you.
 
 # all keys under 'collaborators' corresponds to a specific colaborator name the corresponding dictionary has data_name, data_path pairs.
@@ -8,4 +8,4 @@
 # collaborator_name,data_directory_path
 
 one,1
-two,2
+two,2

openfl-workspace/tf_2dunet/plan/plan.yaml → openfl-workspace/keras_2dunet/plan/plan.yaml

@@ -5,9 +5,9 @@ aggregator :
   defaults : plan/defaults/aggregator.yaml
   template : openfl.component.Aggregator
   settings :
-    init_state_path : save/tf_2dunet_brats_init.pbuf
-    last_state_path : save/tf_2dunet_brats_latest.pbuf
-    best_state_path : save/tf_2dunet_brats_best.pbuf
+    init_state_path : save/init.pbuf
+    last_state_path : save/latest.pbuf
+    best_state_path : save/best.pbuf
     rounds_to_train : 10
     db_store_rounds : 2
 
@@ -20,7 +20,7 @@ collaborator :
 
 data_loader :
   defaults : plan/defaults/data_loader.yaml
-  template : src.tfbrats_inmemory.TensorFlowBratsInMemory
+  template : src.dataloader.KerasBratsInMemory
   settings :
     batch_size: 64
     percent_train: 0.8
@@ -29,7 +29,7 @@ data_loader :
 
 task_runner :
   defaults : plan/defaults/task_runner.yaml
-  template : src.tf_2dunet.TensorFlow2DUNet
+  template : src.taskrunner.Keras2DUNet
 
 network :
   defaults : plan/defaults/network.yaml
@@ -38,7 +38,30 @@ assigner :
   defaults : plan/defaults/assigner.yaml
 
 tasks :
-  defaults : plan/defaults/tasks_tensorflow.yaml
+  aggregated_model_validation:
+    function : validate_task
+    kwargs   :
+      batch_size : 32
+      apply      : global
+      metrics    :
+        - acc
+
+  locally_tuned_model_validation:
+    function : validate_task
+    kwargs   :
+      batch_size : 32
+      apply      : local
+      metrics    :
+        - acc
+
+  train:
+    function : train_task
+    kwargs   :
+      batch_size  : 32
+      metrics     :
+      - loss
+      epochs      : 1
+
 
 compression_pipeline :
-  defaults : plan/defaults/compression_pipeline.yaml
+  defaults : plan/defaults/compression_pipeline.yaml

openfl-workspace/tf_2dunet/requirements.txt → openfl-workspace/keras_2dunet/requirements.txt

@@ -1,3 +1,4 @@
+keras==3.6.0
 nibabel
 setuptools>=65.5.1 # not directly required, pinned by Snyk to avoid a vulnerability
-tensorflow==2.13
+tensorflow==2.18.0

openfl-workspace/tf_2dunet/src/tfbrats_inmemory.py → openfl-workspace/keras_2dunet/src/dataloader.py

@@ -3,12 +3,12 @@
 
 """You may copy this file as the starting point of your own model."""
 
-from openfl.federated import TensorFlowDataLoader
+from openfl.federated import KerasDataLoader
 from .brats_utils import load_from_nifti
 
 
-class TensorFlowBratsInMemory(TensorFlowDataLoader):
-    """TensorFlow Data Loader for the BraTS dataset."""
+class KerasBratsInMemory(KerasDataLoader):
+    """Keras Data Loader for the BraTS dataset."""
 
     def __init__(self, data_path, batch_size, percent_train=0.8, pre_split_shuffle=True, **kwargs):
         """Initialize.

openfl-workspace/keras_2dunet/src/taskrunner.py

@@ -0,0 +1,158 @@
+# 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."""
+
+import tensorflow as tf
+import keras
+
+from openfl.federated import KerasTaskRunner
+
+
+class Keras2DUNet(KerasTaskRunner):
+    """Initialize.
+
+    Args:
+        **kwargs: Additional parameters to pass to the function
+
+    """
+
+    def __init__(self, **kwargs):
+        """Initialize.
+
+        Args:
+            **kwargs: Additional parameters to pass to the function
+
+        """
+        super().__init__(**kwargs)
+
+        self.model = self.build_model(self.data_loader.get_feature_shape(), use_upsampling=True, **kwargs)
+        self.model.summary(print_fn=self.logger.info, line_length=120)
+        self.initialize_tensorkeys_for_functions()
+
+
+    def build_model(self, input_shape,
+                use_upsampling=False,
+                n_cl_out=1,
+                dropout=0.2,
+                activation_function='relu',
+                seed=0xFEEDFACE,
+                depth=5,
+                dropout_at=None,
+                initial_filters=32,
+                batch_norm=True,
+                **kwargs):
+        """Define the TensorFlow model.
+
+        Args:
+            input_shape: input shape of the model
+            use_upsampling (bool): True = use bilinear interpolation;
+                                False = use transposed convolution (Default=False)
+            n_cl_out (int): Number of channels in input layer (Default=1)
+            dropout (float): Dropout percentage (Default=0.2)(Default = True)
+            activation_function: The activation function to use after convolutional layers (Default='relu')
+            seed: random seed (Default=0xFEEDFACE)
+            depth (int): Number of max pooling layers in encoder (Default=5)
+            dropout_at: Layers to perform dropout after (Default=[2,3])
+            initial_filters (int): Number of filters in first convolutional layer (Default=32)
+            batch_norm (bool): True = use batch normalization (Default=True)
+            **kwargs: Additional parameters to pass to the function
+
+        """
+        if dropout_at is None:
+            dropout_at = [2, 3]
+
+        inputs = keras.layers.Input(shape=input_shape, name='Images')
+
+        if activation_function == 'relu':
+            activation = tf.nn.relu
+        elif activation_function == 'leakyrelu':
+            activation = tf.nn.leaky_relu
+
+        params = {
+            'activation': activation,
+            'kernel_initializer': keras.initializers.he_uniform(seed=seed),
+            'kernel_size': (3, 3),
+            'padding': 'same',
+        }
+
+        convb_layers = {}
+
+        net = inputs
+        filters = initial_filters
+        for i in range(depth):
+            name = f'conv{i + 1}a'
+            net = keras.layers.Conv2D(name=name, filters=filters, **params)(net)
+            if i in dropout_at:
+                net = keras.layers.Dropout(dropout)(net)
+            name = f'conv{i + 1}b'
+            net = keras.layers.Conv2D(name=name, filters=filters, **params)(net)
+            if batch_norm:
+                net = keras.layers.BatchNormalization()(net)
+            convb_layers[name] = net
+            # only pool if not last level
+            if i != depth - 1:
+                name = f'pool{i + 1}'
+                net = keras.layers.MaxPooling2D(name=name, pool_size=(2, 2))(net)
+                filters *= 2
+
+        # do the up levels
+        filters //= 2
+        for i in range(depth - 1):
+            if use_upsampling:
+                up = keras.layers.UpSampling2D(
+                    name=f'up{depth + i + 1}', size=(2, 2))(net)
+            else:
+                up = keras.layers.Conv2DTranspose(
+                    name='transConv6', filters=filters,
+                    kernel_size=(2, 2), strides=(2, 2), padding='same')(net)
+            net = keras.layers.concatenate(
+                [up, convb_layers[f'conv{depth - i - 1}b']],
+                axis=-1
+            )
+            net = keras.layers.Conv2D(
+                name=f'conv{depth + i + 1}a',
+                filters=filters, **params)(net)
+            net = keras.layers.Conv2D(
+                name=f'conv{depth + i + 1}b',
+                filters=filters, **params)(net)
+            filters //= 2
+        net = keras.layers.Conv2D(name='Mask', filters=n_cl_out,
+                                    kernel_size=(1, 1),
+                                    activation='sigmoid')(net)
+        model = keras.models.Model(inputs=[inputs], outputs=[net])
+
+
+        self.optimizer = keras.optimizers.RMSprop(1e-2)
+        model.compile(
+            loss=self.dice_coef_loss,
+            optimizer=self.optimizer,
+            metrics=["acc"]
+        )
+
+        return model
+
+    def dice_coef_loss(self, y_true, y_pred, smooth=1.0):
+        """Dice coefficient loss.
+
+        Calculate the -log(Dice Coefficient) loss
+
+        Args:
+            y_true: Ground truth annotation array
+            y_pred: Prediction array from model
+            smooth (float): Laplace smoothing factor (Default=1.0)
+        Returns:
+            float: -log(Dice cofficient) metric
+        """
+        intersection = tf.reduce_sum(y_true * y_pred, axis=(1, 2, 3))
+
+        term1 = -tf.math.log(tf.constant(2.0) * intersection + smooth)
+        term2 = tf.math.log(tf.reduce_sum(y_true, axis=(1, 2, 3))
+                    + tf.reduce_sum(y_pred, axis=(1, 2, 3)) + smooth)
+
+        term1 = tf.reduce_mean(term1)
+        term2 = tf.reduce_mean(term2)
+
+        loss = term1 + term2
+
+        return loss

openfl-workspace/keras_cnn_mnist/plan/plan.yaml

@@ -5,9 +5,9 @@ aggregator :
   defaults : plan/defaults/aggregator.yaml
   template : openfl.component.Aggregator
   settings :
-    init_state_path : save/keras_cnn_mnist_init.pbuf
-    best_state_path : save/keras_cnn_mnist_best.pbuf
-    last_state_path : save/keras_cnn_mnist_last.pbuf
+    init_state_path : save/init.pbuf
+    best_state_path : save/best.pbuf
+    last_state_path : save/last.pbuf
     rounds_to_train : 10
 
 collaborator :
@@ -19,15 +19,15 @@ collaborator :
 
 data_loader :
   defaults : plan/defaults/data_loader.yaml
-  template : src.tfmnist_inmemory.TensorFlowMNISTInMemory
+  template : src.dataloader.KerasMNISTInMemory
   settings :
     collaborator_count : 2
     data_group_name    : mnist
     batch_size         : 256
 
 task_runner :
   defaults : plan/defaults/task_runner.yaml
-  template : src.keras_cnn.KerasCNN
+  template : src.taskrunner.KerasCNN
 
 network :
   defaults : plan/defaults/network.yaml
@@ -40,3 +40,7 @@ tasks :
 
 compression_pipeline :
   defaults : plan/defaults/compression_pipeline.yaml
+  # To use different Compression Pipeline, uncomment the following lines
+  # template : openfl.pipelines.KCPipeline
+  # settings :
+  #   n_clusters : 6

openfl-workspace/keras_cnn_mnist/requirements.txt

@@ -1,2 +1,3 @@
-numpy==1.23.5
-tensorflow==2.13
+keras==3.6.0
+tensorflow==2.18.0
+

openfl-workspace/keras_cnn_mnist/src/tfmnist_inmemory.py → openfl-workspace/keras_cnn_mnist/src/dataloader.py

@@ -3,12 +3,12 @@
 
 """You may copy this file as the starting point of your own model."""
 
-from openfl.federated import TensorFlowDataLoader
+from openfl.federated import KerasDataLoader
 from .mnist_utils import load_mnist_shard
 
 
-class TensorFlowMNISTInMemory(TensorFlowDataLoader):
-    """TensorFlow Data Loader for MNIST Dataset."""
+class KerasMNISTInMemory(KerasDataLoader):
+    """Data Loader for MNIST Dataset."""
 
     def __init__(self, data_path, batch_size, **kwargs):
         """

openfl-workspace/keras_cnn_mnist/src/keras_cnn.py → openfl-workspace/keras_cnn_mnist/src/taskrunner.py

@@ -1,13 +1,12 @@
-# Copyright (C) 2020-2021 Intel Corporation
+# Copyright (C) 2020-2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
 """You may copy this file as the starting point of your own model."""
 
-import tensorflow.keras as ke
-from tensorflow.keras import Sequential
-from tensorflow.keras.layers import Conv2D
-from tensorflow.keras.layers import Dense
-from tensorflow.keras.layers import Flatten
+from keras.models import Sequential
+from keras.layers import Conv2D
+from keras.layers import Dense
+from keras.layers import Flatten
 
 from openfl.federated import KerasTaskRunner
 
@@ -50,7 +49,7 @@ def build_model(self,
             num_classes (int): The number of classes of the dataset
 
         Returns:
-            tensorflow.python.keras.engine.sequential.Sequential: The model defined in Keras
+            keras.models.Sequential: The model defined in Keras
 
         """
         model = Sequential()
@@ -72,14 +71,8 @@ def build_model(self,
 
         model.add(Dense(num_classes, activation='softmax'))
 
-        model.compile(loss=ke.losses.categorical_crossentropy,
-                      optimizer=ke.optimizers.legacy.Adam(),
-                      metrics=['accuracy'])
-
-        # initialize the optimizer variables
-        opt_vars = model.optimizer.variables()
-
-        for v in opt_vars:
-            v.initializer.run(session=self.sess)
+        model.compile(loss="categorical_crossentropy",
+                      optimizer="adam",
+                      metrics=["accuracy"])
 
         return model