Support Saving Tensors in Graph Mode with add_for_mode (aws#353)

Author: NihalHarish

smdebug/tensorflow/keras.py

@@ -350,6 +350,15 @@ def _prepare_layers(self, mode):
             for w in weights:
                 self._check_and_add_layer_tensor(mode, layer, "weight", w)
 
+    def _prepare_non_layer_tensors(self):
+        for coll in self.collection_manager.get_collections().values():
+            collection_values = coll.get_tensors()
+            for tensor_ref in collection_values:
+                if tensor_ref.name not in self.tensor_to_collections:
+                    self.tensor_to_collections[tensor_ref.name] = {coll}
+                elif coll not in self.tensor_to_collections[tensor_ref.name]:
+                    self.tensor_to_collections[tensor_ref.name].add(coll)
+
     def _prepare_tensors_available_post_step(self):
         # for gradients, optimizer_variables
         custom_collections, _ = self._get_custom_and_default_collections()
@@ -359,7 +368,8 @@ def _prepare_tensors_available_post_step(self):
             self.get_collection(name=CollectionKeys.OUTPUTS),
             self.get_collection(name=CollectionKeys.INPUTS),
         ]:
-            for tensor_ref in coll.get_tensors():
+            collection_values = coll.get_tensors()
+            for tensor_ref in collection_values:
                 if tensor_ref.name not in self.tensor_to_collections:
                     self.tensor_to_collections[tensor_ref.name] = {coll}
                 elif coll not in self.tensor_to_collections[tensor_ref.name]:
@@ -729,6 +739,7 @@ def _on_any_batch_begin(self, batch, mode, logs=None):
                 self._get_exec_function(mode)
             ):
                 self._prepare_layers(mode)
+                self._prepare_non_layer_tensors()
                 self._prepare_tensors_available_post_step()
                 self._prepared_tensors[mode] = True
                 # below should be after tensors are processed,

tests/tensorflow2/test_functional_model.py

@@ -0,0 +1,67 @@
+# Third Party
+import tensorflow as tf
+from tensorflow.keras.layers import Conv2D, Dense, Dropout, Flatten, Input, MaxPooling2D
+
+# First Party
+import smdebug.tensorflow as smd
+
+
+def create_dataset():
+    # Download and load MNIST dataset.
+    (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data("MNIST-data")
+    x_train, x_test = x_train / 255.0, x_test / 255.0
+
+    # Add a channels dimension
+    x_train = x_train[..., tf.newaxis]
+    x_test = x_test[..., tf.newaxis]
+
+    train_ds = (
+        tf.data.Dataset.from_tensor_slices((x_train, y_train)).shuffle(10000, seed=123).batch(2)
+    )
+    test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(2)
+
+    return train_ds, test_ds
+
+
+def test_functional_model(out_dir, tf_eager_mode):
+    if tf_eager_mode is False:
+        tf.compat.v1.disable_eager_execution()
+    else:
+        return
+    num_classes = 10
+    train_ds, test_ds = create_dataset()
+
+    # Input image dimensions
+    img_rows, img_cols = 28, 28
+
+    img_inputs = Input(shape=(28, 28, 1))
+    x = Conv2D(32, kernel_size=(3, 3), activation="relu")(img_inputs)
+    x1 = Conv2D(64, (3, 3), activation="relu")(x)
+    x = MaxPooling2D(pool_size=(2, 2))(x1)
+    x = Dropout(0.25)(x)
+    x = Flatten()(x)
+    x = Dense(128, activation="relu")(x)
+    x = Dropout(0.5)(x)
+    out = Dense(num_classes, activation="softmax")(x)
+
+    model = tf.keras.models.Model(inputs=img_inputs, outputs=out)
+
+    smd_callback = smd.KerasHook(
+        export_tensorboard=False, out_dir=out_dir, include_collections=["custom"]
+    )
+
+    smd_callback.get_collection("custom").add_for_mode([x1], mode=smd.modes.TRAIN)
+    smd_callback.save_config = smd.SaveConfig(save_interval=1)
+    opt = tf.keras.optimizers.Adadelta(1.0)
+
+    model.compile(
+        loss=tf.keras.losses.sparse_categorical_crossentropy,
+        optimizer=opt,
+        experimental_run_tf_function=False,
+    )
+
+    callbacks = [smd_callback]
+    model.fit(train_ds, epochs=1, steps_per_epoch=100, callbacks=callbacks)
+
+    trial = smd.create_trial(out_dir)
+    assert len(trial.tensor_names(collection="custom")) == 1