Reduce data logged to tensorboard

Summary:
We were logging learning rate and loss for every single step, which makes
tensorboard too slow to load in long training runs. Log every 10th step, which
should be enough for all cases: we always log at the end of every phase as
well.

Differential Revision: D20441202

fbshipit-source-id: 6516dee931dcc114b1ccd3074bf4f0fbbd61719f

Author: vreis

classy_vision/hooks/tensorboard_plot_hook.py

@@ -52,13 +52,15 @@ def __init__(self, tb_writer) -> None:
         self.tb_writer = tb_writer
         self.learning_rates: Optional[List[float]] = None
         self.wall_times: Optional[List[float]] = None
-        self.num_steps_global: Optional[List[int]] = None
+        self.num_updates: Optional[List[int]] = None
+        self.log_period = 10
 
     def on_phase_start(self, task: "tasks.ClassyTask") -> None:
         """Initialize losses and learning_rates."""
         self.learning_rates = []
         self.wall_times = []
-        self.num_steps_global = []
+        self.num_updates = []
+        self.step_idx = 0
 
     def on_step(self, task: "tasks.ClassyTask") -> None:
         """Store the observed learning rates."""
@@ -70,11 +72,14 @@ def on_step(self, task: "tasks.ClassyTask") -> None:
             # Only need to log the average loss during the test phase
             return
 
-        learning_rate_val = task.optimizer.parameters.lr
+        if self.step_idx % self.log_period == 0:
+            learning_rate_val = task.optimizer.parameters.lr
 
-        self.learning_rates.append(learning_rate_val)
-        self.wall_times.append(time.time())
-        self.num_steps_global.append(task.num_updates)
+            self.learning_rates.append(learning_rate_val)
+            self.wall_times.append(time.time())
+            self.num_updates.append(task.num_updates)
+
+        self.step_idx += 1
 
     def on_phase_end(self, task: "tasks.ClassyTask") -> None:
         """Add the losses and learning rates to tensorboard."""
@@ -97,7 +102,7 @@ def on_phase_end(self, task: "tasks.ClassyTask") -> None:
 
         if task.train:
             for loss, learning_rate, global_step, wall_time in zip(
-                task.losses, self.learning_rates, self.num_steps_global, self.wall_times
+                task.losses, self.learning_rates, self.num_updates, self.wall_times
             ):
                 loss /= task.get_batchsize_per_replica()
                 self.tb_writer.add_scalar(

test/manual/hooks_tensorboard_plot_hook_test.py

@@ -146,6 +146,7 @@ def flush(self):
 
         writer = DummySummaryWriter()
         hook = TensorboardPlotHook(writer)
+        hook.log_period = 1
         task.set_hooks([hook])
         task.optimizer.param_schedulers["lr"] = mock_lr_scheduler