removed the VisualizeCallback change

Author: saltytine

basic_pitch/callbacks.py

@@ -17,7 +17,7 @@
 
 import os
 
-from typing import Any, Dict, Callable
+from typing import Any, Dict
 
 import tensorflow as tf
 
@@ -27,6 +27,7 @@
 
 
 class VisualizeCallback(tf.keras.callbacks.Callback):
+    # TODO RACHEL make this WAY faster
     """
     Callback to run during training to create tensorboard visualizations per epoch.
 
@@ -45,60 +46,32 @@ def __init__(
         tensorboard_dir: str,
         sonify: bool,
         contours: bool,
-        max_batches: int = 2,
-        prefetch_batches: int = 2,
-        use_tf_function: bool = True,
     ):
         super().__init__()
-        self.train_ds = train_ds.take(max_batches).prefetch(prefetch_batches)
-        self.validation_ds = validation_ds.take(max_batches).prefetch(prefetch_batches)
+        self.train_iter = iter(train_ds)
+        self.validation_iter = iter(validation_ds)
         self.tensorboard_dir = os.path.join(tensorboard_dir, TENSORBOARD_LOGS_SUBDIR)
-        self.file_writer = tf.summary.create_file_writer(self.tensorboard_dir)
+        self.file_writer = tf.summary.create_file_writer(tensorboard_dir)
         self.sonify = sonify
         self.contours = contours
-        self.use_tf_function = use_tf_function
-
-        self.train_iter = iter(self.train_ds)
-        self.validation_iter = iter(self.validation_ds)
-
-        self._predict_fn = None
-
-    def set_model(self, model: tf.keras.Model) -> None:
-        super().set_model(model)
-        if self.use_tf_function:
-            def fast_predict(inputs: tf.Tensor) -> Any:
-                return model(inputs, training=False)
-            self._predict_fn = tf.function(fast_predict)
-        else:
-            self._predict_fn = model.predict
-
-    def _predict(self, inputs: tf.Tensor) -> Any:
-        if self._predict_fn is not None:
-            outputs = self._predict_fn(inputs)
-            if isinstance(outputs, dict):
-                outputs = {k: v.numpy() if hasattr(v, "numpy") else v for k, v in outputs.items()}
-            return outputs
-        else:
-            return self.model.predict(inputs)
 
     def on_epoch_end(self, epoch: int, logs: Dict[Any, Any]) -> None:
-        for stage, ds, loss_key in [
-            ("train", self.train_ds, "loss"),
-            ("validation", self.validation_ds, "val_loss"),
+        # the first two outputs of generator needs to be the input and the targets
+        train_inputs, train_targets = next(self.train_iter)[:2]
+        validation_inputs, validation_targets = next(self.validation_iter)[:2]
+        for stage, inputs, targets, loss in [
+            ("train", train_inputs, train_targets, logs["loss"]),
+            ("validation", validation_inputs, validation_targets, logs["val_loss"]),
         ]:
-            for batch in ds:
-                inputs, targets = batch[:2]
-                outputs = self._predict(inputs)
-                loss_val = logs.get(loss_key)
-                visualize.visualize_transcription(
-                    self.file_writer,
-                    stage,
-                    inputs,
-                    targets,
-                    outputs,
-                    float(loss_val) if loss_val is not None else 0.0,
-                    epoch,
-                    sonify=self.sonify,
-                    contours=self.contours,
-                )
-                break
+            outputs = self.model.predict(inputs)
+            visualize.visualize_transcription(
+                self.file_writer,
+                stage,
+                inputs,
+                targets,
+                outputs,
+                loss,
+                epoch,
+                sonify=self.sonify,
+                contours=self.contours,
+            )