Repository: NVIDIA/NVFlare

Issue:
Add a TensorBoard metric streaming example
Show how to stream metrics during training to the server and create central tensorboard event files.


Discussion:

Resolving patch:
diff --git a/examples/cifar10/pt/learners/cifar10_learner.py b/examples/cifar10/pt/learners/cifar10_learner.py
--- a/examples/cifar10/pt/learners/cifar10_learner.py
+++ b/examples/cifar10/pt/learners/cifar10_learner.py
@@ -31,7 +31,6 @@
 from nvflare.app_common.abstract.learner_spec import Learner
 from nvflare.app_common.app_constant import AppConstants, ModelName, ValidateType
 from nvflare.app_common.pt.pt_fedproxloss import PTFedProxLoss
-from nvflare.app_common.widgets.streaming import send_analytic_dxo, write_scalar
 
 
 class CIFAR10Learner(Learner):
@@ -44,7 +43,7 @@ def __init__(
         lr: float = 1e-2,
         fedproxloss_mu: float = 0.0,
         central: bool = False,
-        stream_tb: bool = False
+        analytic_sender_id: str = "analytic_sender",
     ):
         """Simple CIFAR-10 Trainer.
 
@@ -53,7 +52,7 @@ def __init__(
             aggregation_epochs: the number of training epochs for a round. Defaults to 1.
             train_task_name: name of the task to train the model.
             submit_model_task_name: name of the task to submit the best local model.
-            stream_tb: whether to send TensorBoard values. Defaults to False.
+            analytic_sender_id: id of `AnalyticsSender` if configured as a client component. If configured, TensorBoard events will be fired. Defaults to "analytic_sender".
 
         Returns:
             a Shareable with the updated local model after running `execute()`
@@ -72,7 +71,7 @@ def __init__(
         self.central = central
 
         self.writer = None
-        self.stream_tb = stream_tb
+        self.analytic_sender_id = analytic_sender_id
 
         # Epoch counter
         self.epoch_of_start_time = 0
@@ -93,8 +92,10 @@ def initialize(self, parts: dict, fl_ctx: FLContext):
         self.local_model_file = os.path.join(self.app_root, "local_model.pt")
         self.best_local_model_file = os.path.join(self.app_root, "best_local_model.pt")
 
-        # Set local tensorboard writer - to be replaced by event
-        self.writer = SummaryWriter(self.app_root)
+        # Select local TensorBoard writer or event-based writer for streaming
+        self.writer = parts.get(self.analytic_sender_id)  # user configured config_fed_client.json for streaming
+        if not self.writer:  # use local TensorBoard writer only
+            self.writer = SummaryWriter(self.app_root)
 
         # Set datalist, here the path and filename are hard-coded, can also be fed as an argument
         site_idx_file_name = os.path.join(self.dataset_root, self.client_id + ".npy")
@@ -164,10 +165,6 @@ def finalize(self, fl_ctx: FLContext):
         # collect threads, close files here
         pass
 
-    def send_tb_scalar(self, tag, number, r, fl_ctx):
-        dxo = write_scalar(tag, number, global_step=r)
-        send_analytic_dxo(comp=self, dxo=dxo, fl_ctx=fl_ctx)
-
     def local_train(self, fl_ctx, train_loader, model_global, abort_signal: Signal, val_freq: int = 0):
         for epoch in range(self.aggregation_epochs):
             if abort_signal.triggered:
@@ -195,8 +192,6 @@ def local_train(self, fl_ctx, train_loader, model_global, abort_signal: Signal,
                 self.optimizer.step()
                 current_step = epoch_len * self.epoch_global + i
                 self.writer.add_scalar("train_loss", loss.item(), current_step)
-                if self.stream_tb:
-                    self.send_tb_scalar("train_loss", loss.item(), current_step, fl_ctx)
             if val_freq > 0 and epoch % val_freq == 0:
                 acc = self.local_valid(self.valid_loader, abort_signal, tb_id="val_acc_local_model", fl_ctx=fl_ctx)
                 if acc > self.best_acc:
@@ -332,8 +327,6 @@ def local_valid(self, valid_loader, abort_signal: Signal, tb_id=None, fl_ctx=Non
             metric = correct / float(total)
             if tb_id:
                 self.writer.add_scalar(tb_id, metric, self.epoch_global)
-                if self.stream_tb:
-                    self.send_tb_scalar(tb_id, metric, self.epoch_global, fl_ctx)
         return metric
 
     def validate(self, shareable: Shareable, fl_ctx: FLContext, abort_signal: Signal) -> Shareable:


Test patch: