chore: make tensorboard logs lighter (#687)

pyannote · Jun 22, 2021 · 2ecdc23 · 2ecdc23
1 parent 03484f0
commit 2ecdc23
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 7 deletions.
diff --git a/pyannote/audio/tasks/segmentation/mixins.py b/pyannote/audio/tasks/segmentation/mixins.py
@@ -495,8 +495,12 @@ def validation_step(self, batch, batch_idx: int):
             logger=True,
         )
 
-        # log first batch visualization every 10 epochs.
-        if self.model.current_epoch % 10 > 0 or batch_idx > 0:
+        # log first batch visualization every 2^n epochs.
+        if (
+            self.model.current_epoch == 0
+            or math.log2(self.model.current_epoch) % 1 > 0
+            or batch_idx > 0
+        ):
             return
 
         # visualize first 9 validation samples of first batch in Tensorboard

diff --git a/pyannote/audio/tasks/segmentation/segmentation.py b/pyannote/audio/tasks/segmentation/segmentation.py
@@ -337,7 +337,7 @@ def training_step(self, batch, batch_idx: int):
         self.model.log(
             f"{self.ACRONYM}@train_seg_loss",
             seg_loss,
-            on_step=True,
+            on_step=False,
             on_epoch=True,
             prog_bar=False,
             logger=True,
@@ -354,7 +354,7 @@ def training_step(self, batch, batch_idx: int):
             self.model.log(
                 f"{self.ACRONYM}@train_vad_loss",
                 vad_loss,
-                on_step=True,
+                on_step=False,
                 on_epoch=True,
                 prog_bar=False,
                 logger=True,
@@ -365,7 +365,7 @@ def training_step(self, batch, batch_idx: int):
         self.model.log(
             f"{self.ACRONYM}@train_loss",
             loss,
-            on_step=True,
+            on_step=False,
             on_epoch=True,
             prog_bar=True,
             logger=True,
@@ -414,8 +414,12 @@ def validation_step(self, batch, batch_idx: int):
             logger=True,
         )
 
-        # log first batch visualization every 10 epochs.
-        if self.model.current_epoch % 10 > 0 or batch_idx > 0:
+        # log first batch visualization every 2^n epochs.
+        if (
+            self.model.current_epoch == 0
+            or math.log2(self.model.current_epoch) % 1 > 0
+            or batch_idx > 0
+        ):
             return
 
         # visualize first 9 validation samples of first batch in Tensorboard