Support Inputs and Labels in the dict format

smdebug/tensorflow/keras.py

@@ -5,6 +5,7 @@
 import tensorflow.compat.v1 as tf
 from tensorflow.python.distribute import values
 from tensorflow.python.framework.indexed_slices import IndexedSlices
+from tensorflow.python.util import nest
 
 # First Party
 from smdebug.core.modes import ModeKeys
@@ -497,6 +498,10 @@ def save_smdebug_logs(self, logs):
                         else set()
                     )
                 for t_name, t_value in tensors_to_save:
+                    if isinstance(t_value, dict):
+                        # flatten the inputs and labels
+                        # since we cannot convert dicts into numpy
+                        t_value = nest.flatten(t_value)
                     self._save_tensor_to_file(t_name, t_value, collections_to_write)
 
     def _save_metrics(self, batch, logs, force_save=False):

tests/tensorflow2/test_support_dicts.py

@@ -0,0 +1,42 @@
+# Third Party
+import numpy as np
+import tensorflow as tf
+
+# First Party
+import smdebug.tensorflow as smd
+from smdebug.core.collection import CollectionKeys
+from smdebug.trials import create_trial
+
+
+def get_data():
+    images = np.zeros((64, 224))
+    labels = np.zeros((64, 5))
+    inputs = {"Image_input": images}
+    outputs = {"output-softmax": labels}
+    return inputs, outputs
+
+
+def create_hook(trial_dir):
+    hook = smd.KerasHook(trial_dir, save_all=True)
+    return hook
+
+
+def create_model():
+    input_layer = tf.keras.layers.Input(name="Image_input", shape=(224), dtype="float32")
+    model = tf.keras.layers.Dense(5)(input_layer)
+    model = tf.keras.layers.Activation("softmax", name="output-softmax")(model)
+    model = tf.keras.models.Model(inputs=input_layer, outputs=[model])
+    return model
+
+
+def test_support_dicts(out_dir):
+    model = create_model()
+    optimizer = tf.keras.optimizers.Adadelta(lr=1.0, rho=0.95, epsilon=None, decay=0.0)
+    model.compile(loss="categorical_crossentropy", optimizer=optimizer)
+    inputs, labels = get_data()
+    smdebug_hook = create_hook(out_dir)
+    model.fit(inputs, labels, batch_size=16, epochs=10, callbacks=[smdebug_hook])
+    model.save(out_dir, save_format="tf")
+    trial = create_trial(out_dir)
+    assert trial.tensor_names(collection=CollectionKeys.INPUTS) == ["model_input"]
+    assert trial.tensor_names(collection=CollectionKeys.OUTPUTS) == ["labels", "predictions"]