From 103127e82c3829532b32ea431795e37055b0f598 Mon Sep 17 00:00:00 2001 From: Mayank Gulati Date: Wed, 18 Aug 2021 14:40:29 +0200 Subject: [PATCH] Tensorboard Strategy Decorator. (#789) Co-authored-by: Mayank Gulati Co-authored-by: Taner Topal --- .gitignore | 1 + examples/simulation/simulation.py | 1 - src/py/flwr/server/utils/__init__.py | 21 +++ src/py/flwr/server/utils/tensorboard.py | 142 +++++++++++++++++++ src/py/flwr/server/utils/tensorboard_test.py | 31 ++++ 5 files changed, 195 insertions(+), 1 deletion(-) create mode 100644 src/py/flwr/server/utils/__init__.py create mode 100644 src/py/flwr/server/utils/tensorboard.py create mode 100644 src/py/flwr/server/utils/tensorboard_test.py diff --git a/.gitignore b/.gitignore index 76564c777a42..34bbb0503162 100644 --- a/.gitignore +++ b/.gitignore @@ -3,6 +3,7 @@ data/ doc/source/api_documentation doc/source/_build +flwr_logs # Editor .vscode diff --git a/examples/simulation/simulation.py b/examples/simulation/simulation.py index 148f7970bdac..37d8b9c6842d 100644 --- a/examples/simulation/simulation.py +++ b/examples/simulation/simulation.py @@ -7,7 +7,6 @@ import numpy as np import tensorflow as tf from flwr.server.strategy import FedAvg - import dataset # Make TensorFlow log less verbose diff --git a/src/py/flwr/server/utils/__init__.py b/src/py/flwr/server/utils/__init__.py new file mode 100644 index 000000000000..9ecf77ca0ab2 --- /dev/null +++ b/src/py/flwr/server/utils/__init__.py @@ -0,0 +1,21 @@ +# Copyright 2020 Adap GmbH. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Contains server side utilities to be used in combination with various +components.""" + + +from .tensorboard import tensorboard as tensorboard + +__all__ = ["tensorboard"] diff --git a/src/py/flwr/server/utils/tensorboard.py b/src/py/flwr/server/utils/tensorboard.py new file mode 100644 index 000000000000..e4a0b09103de --- /dev/null +++ b/src/py/flwr/server/utils/tensorboard.py @@ -0,0 +1,142 @@ +# Copyright 2020 Adap GmbH. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Flower TensorBoard utilities.""" +import os +from datetime import datetime +from typing import Callable, Dict, List, Optional, Tuple, TypeVar, cast + +try: + import tensorflow as tf +except ImportError: + tf = None + +from flwr.common import EvaluateRes, Scalar +from flwr.server.client_proxy import ClientProxy +from flwr.server.strategy import Strategy + +TBW = TypeVar("TBW") + + +def tensorboard(logdir: str) -> Callable[[Strategy], TBW]: + """TensorBoard logger for Flower strategies. + + It will log loss, num_examples and all metrics which are of type float or int. + + This can either be used as a decorator as shown in the example variant 1 + or directly as shown in the example variant 2. + + Example: + # Variant 1 + @tensorboard(logdir=LOGDIR) + class CustomStrategy(FedAvg): + pass + + strategy = CustomStrategy() + + # Variant 2 + strategy = tensorboard(logdir=LOGDIR)(FedAvg)() + """ + print( + "\n\t\033[32mStart TensorBoard with the following parameters" + + f"\n\t$ tensorboard --logdir {logdir}\033[39m\n" + ) + # Create logdir if it does not yet exist + os.makedirs(logdir, exist_ok=True) + + # To allow multiple runs and group those we will create a subdir + # in the logdir which is named as number of directories in logdir + 1 + run_id = str( + len( + [ + name + for name in os.listdir(logdir) + if os.path.isdir(os.path.join(logdir, name)) + ] + ) + ) + run_id = run_id + "-" + datetime.now().strftime("%Y%m%dT%H%M%S") + logdir_run = os.path.join(logdir, run_id) + + def decorator(strategy_class: Strategy) -> TBW: + """Return overloaded Strategy Wrapper.""" + + class TBWrapper(strategy_class): # type: ignore + """Strategy wrapper which hooks into some methods for TensorBoard + logging.""" + + def aggregate_evaluate( + self, + rnd: int, + results: List[Tuple[ClientProxy, EvaluateRes]], + failures: List[BaseException], + ) -> Tuple[Optional[float], Dict[str, Scalar]]: + """Hooks into aggregate_evaluate for TensorBoard logging + purpose.""" + # Execute decorated function and extract results for logging + # They will be returned at the end of this function but also + # used for logging + loss_aggregated, config = super().aggregate_evaluate( + rnd, + results, + failures, + ) + + if rnd < 0: + # rnd < 0 is currently planned to be removed and should not be + # used but rather with backwards compatibility in mind ignored + return loss_aggregated, config + + # Server logs + writer = tf.summary.create_file_writer( + os.path.join(logdir_run, "server") + ) + + # Write aggregated loss + with writer.as_default(step=rnd): # pylint: disable=not-context-manager + tf.summary.scalar( + "server/loss_aggregated", loss_aggregated, step=rnd + ) + writer.flush() + + if len(results) == 0: + return loss_aggregated, config + + # Client logs + for client, evaluate_res in results: + loss, num_examples, metrics = ( + evaluate_res.loss, + evaluate_res.num_examples, + evaluate_res.metrics, + ) + + writer = tf.summary.create_file_writer( + os.path.join(logdir_run, "clients", client.cid) + ) + with writer.as_default( # pylint: disable=not-context-manager + step=rnd + ): + tf.summary.scalar("clients/loss", loss) + tf.summary.scalar("clients/num_examples", num_examples) + if metrics is not None: + for key, value in metrics.items(): + if type(value) in [int, float]: + tf.summary.scalar(f"clients/{key}", value) + writer.flush() + + return loss_aggregated, config + + return cast(TBW, TBWrapper) + + return decorator diff --git a/src/py/flwr/server/utils/tensorboard_test.py b/src/py/flwr/server/utils/tensorboard_test.py new file mode 100644 index 000000000000..feb15112f09d --- /dev/null +++ b/src/py/flwr/server/utils/tensorboard_test.py @@ -0,0 +1,31 @@ +# Copyright 2020 Adap GmbH. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""(De-)Serialization Tests.""" + + +from tempfile import TemporaryDirectory + +from .tensorboard import tensorboard + + +def test_tensorboard() -> None: + """Test if tensorboard returns a decorator.""" + + # Prepare + with TemporaryDirectory() as tmpdir: + my_decorator = tensorboard(tmpdir) + + # Assert + assert callable(my_decorator)