Save Optimizer Variables With Keras Fit API In Eager Mode (aws#218)

Author: NihalHarish

smdebug/tensorflow/base_hook.py

@@ -399,7 +399,11 @@ def _get_collections_with_tensor(self, tf_tensor_name) -> Set["Collection"]:
         # tensors are not matched with collections at preparation time.
         # Call core/hook.py's _get_collections_with_tensor() where tensors are
         # matched with collections by regex
-        if self.tape:
+        if self.tape or (
+            tf_tensor_name not in self.tensor_to_collections
+            and is_tf_version_2x()
+            and tf.executing_eagerly()
+        ):
             return super()._get_collections_with_tensor(tf_tensor_name)
         return self.tensor_to_collections[tf_tensor_name]
 
@@ -457,8 +461,8 @@ def set_gradients(self, gradients=None, gradients_and_variables=None):
         :param gradients_and_variables: list of tuples [(tf.Tensor/tf.Variable, tf.Tensor/tf.Variable)...]
             list of tuples representing gradients and weights
         """
-        # TF 2.x doesn't provide gradient/optimizer variable names and values by default.
-        # Skipping set_gradients and set_optimizer_variables for Tf 2.x until there is
+        # TF 2.x provides only symbolic gradient variables that do not provide access to their values.
+        # Skipping set_gradients for Tf 2.x until there is
         # support to pass names and values from TF side.
 
         # From TF 2.2, executing_eagerly_outside_functions() can be used as
@@ -482,18 +486,12 @@ def set_optimizer_variables(self, optimizer_variables):
         This method helps find the optimizer variables (such as momentum)
         :param optimizer_variables: list of tf.Variables/tf.Tensors/tf.MirroredVariables
         """
-        # TF 2.x doesn't provide gradient/optimizer variable names and values by default.
-        # Skipping set_gradients and set_optimizer_variables for Tf 2.x until there is
-        # support to pass names and values from TF side.
-
         # From TF 2.2, executing_eagerly_outside_functions() can be used as
         # ops.executing_eagerly_outside_functions() or tf.compat.v1.executing_eagerly_outside_functions().
         # But in TF 2.1, only ops.executing_eagerly_outside_functions() is valid
-        if is_tf_version_2x() and ops.executing_eagerly_outside_functions():
-            return
         # since this is done for each variable at a time for keras, not checking if set already
         self.collection_manager.get(CollectionKeys.OPTIMIZER_VARIABLES).add_for_mode(
-            optimizer_variables, ModeKeys.TRAIN
+            optimizer_variables, self.mode
         )
 
     @staticmethod

smdebug/tensorflow/keras.py

@@ -6,6 +6,7 @@
 from tensorflow.python.distribute import values
 
 # First Party
+from smdebug.core.collection import DEFAULT_TF_COLLECTIONS
 from smdebug.core.modes import ModeKeys
 from smdebug.core.utils import match_inc
 from smdebug.tensorflow.callable_cache import CallableCache
@@ -317,12 +318,28 @@ def _prepare_layers(self, mode):
 
     def _prepare_non_layer_tensors(self):
         # for gradients, optimizer_variables
+        custom_collections = set()
+        default_tf_collection = set()
+
         for coll in self.collection_manager.get_collections().values():
-            for tensor_ref in coll.get_tensors():
+            if coll.name not in DEFAULT_TF_COLLECTIONS:
+                custom_collections.add(coll)
+            else:
+                default_tf_collection.add(coll)
+
+        for default_coll in default_tf_collection:
+            for tensor_ref in default_coll.get_tensors():
                 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)
+                    self.tensor_to_collections[tensor_ref.name] = {default_coll}
+                elif default_coll not in self.tensor_to_collections[tensor_ref.name]:
+                    self.tensor_to_collections[tensor_ref.name].add(default_coll)
+
+                # Add tensor to custom collections
+                for custom_coll in custom_collections:
+                    if match_inc(tensor_ref.name, custom_coll.include_regex):
+                        custom_coll.add_for_mode(tensor_ref.tf_obj, self.mode)
+                        if custom_coll not in self.tensor_to_collections[tensor_ref.name]:
+                            self.tensor_to_collections[tensor_ref.name].add(custom_coll)
 
     def _prepare_tensors_for_step(self, mode):
         self.tensor_refs_to_save_this_step = set()
@@ -550,12 +567,32 @@ def on_test_batch_begin(self, batch, logs=None):
     def on_predict_batch_begin(self, batch, logs=None):
         self._on_any_batch_begin(batch, ModeKeys.PREDICT, logs=logs)
 
+    def _write_optimizer_variables(self):
+        optimizer_collections = self.collection_manager.get(CollectionKeys.OPTIMIZER_VARIABLES)
+        collections = self._get_collections_to_save_for_step()
+        for tensor_ref in optimizer_collections.get_tensors(self.mode):
+            for coll in collections:
+                if coll in self.tensor_to_collections[tensor_ref.name]:
+                    tensor = tensor_ref.tf_obj
+                    self._save_for_tensor(
+                        tensor_name=tensor.name,
+                        tensor_value=tensor.value(),
+                        check_before_write=False,
+                    )
+
     def _on_any_batch_end(self, batch, mode, logs=None):
         if self._is_not_supported():
             return
 
         if not is_tf_version_2x() or (is_tf_version_2x() and not tf.executing_eagerly()):
             self._remove_fetches_and_callbacks(mode)
+
+        if is_tf_version_2x() and tf.executing_eagerly():
+            # Need to prepare non layer tensors again since
+            # some tensors only become available on  batch end
+            self._prepare_non_layer_tensors()
+            self._write_optimizer_variables()
+
         self._save_tensors_post_step(batch, logs)
 
         if self._prepared_tensors[mode]:

tests/tensorflow/keras/test_keras.py

@@ -340,7 +340,7 @@ def test_include_regex(out_dir):
     tr = create_trial_fast_refresh(out_dir)
     tnames = tr.tensor_names(collection="custom_coll")
 
-    assert len(tnames) == 8
+    assert len(tnames) == 12
     for tname in tnames:
         assert tr.tensor(tname).value(0) is not None
 

tests/tensorflow/keras/test_keras_mirrored.py

@@ -502,7 +502,7 @@ def test_include_regex(out_dir):
     tr = create_trial_fast_refresh(out_dir)
     tnames = tr.tensor_names(collection="custom_coll")
 
-    assert len(tnames) == 4 + 3 * strategy.num_replicas_in_sync
+    assert len(tnames) == 4 + 4 + 3 * strategy.num_replicas_in_sync
     for tname in tnames:
         assert tr.tensor(tname).value(0) is not None
 

tests/tensorflow2/test_keras.py

@@ -17,6 +17,7 @@
 from smdebug.core.access_layer import has_training_ended
 from smdebug.core.collection import CollectionKeys
 from smdebug.core.json_config import CONFIG_FILE_PATH_ENV_STR
+from smdebug.core.modes import ModeKeys
 from smdebug.core.reduction_config import ALLOWED_NORMS, ALLOWED_REDUCTIONS
 from smdebug.exceptions import TensorUnavailableForStep
 from smdebug.tensorflow import ReductionConfig, SaveConfig
@@ -395,11 +396,21 @@ def test_keras_fit(out_dir, tf_eager_mode, saveall):
     # can't save gradients in TF 2.x eager mode
     if saveall:  # save losses, metrics, weights, biases
         if tf_eager_mode:
-            assert len(trial.tensor_names()) == (7 if is_tf_2_2() else 8)
+            assert len(trial.tensor_names()) == (12 if is_tf_2_2() else 13)
         else:
             assert len(trial.tensor_names()) == 21
         assert len(trial.tensor_names(collection=CollectionKeys.BIASES)) == 2
         assert len(trial.tensor_names(collection=CollectionKeys.WEIGHTS)) == 2
+        assert len(trial.tensor_names(collection=CollectionKeys.OPTIMIZER_VARIABLES)) == 5
+        assert (
+            len(
+                trial.tensor_names(
+                    collection=CollectionKeys.OPTIMIZER_VARIABLES, mode=ModeKeys.EVAL
+                )
+            )
+            == 0,
+            "No Optimizer Variables Should be Saved in EVAL Mode",
+        )
     else:  # save the default losses and metrics
         assert len(trial.tensor_names()) == (3 if is_tf_2_2() and tf_eager_mode else 4)
     assert len(trial.tensor_names(collection=CollectionKeys.LOSSES)) == 1
@@ -540,6 +551,7 @@ def test_include_collections(out_dir, tf_eager_mode):
         CollectionKeys.OUTPUTS,
         CollectionKeys.METRICS,
         CollectionKeys.OPTIMIZER_VARIABLES,
+        "custom_optimizer_variables",
     ]
     save_config = SaveConfig(save_interval=3)
     hook = smd.KerasHook(
@@ -548,16 +560,18 @@ def test_include_collections(out_dir, tf_eager_mode):
         include_collections=include_collections,
         reduction_config=ReductionConfig(norms=ALLOWED_NORMS, reductions=ALLOWED_REDUCTIONS),
     )
+    hook.get_collection("custom_optimizer_variables").include("Adam")
     helper_keras_fit(out_dir, hook=hook, steps=["train", "eval", "predict"], eager=tf_eager_mode)
 
     trial = smd.create_trial(path=out_dir)
     # can't save gradients in TF 2.x
     if tf_eager_mode:
-        assert len(trial.tensor_names()) == (7 if is_tf_2_2() else 8)
+        assert len(trial.tensor_names()) == (12 if is_tf_2_2() else 13)
     else:
         assert len(trial.tensor_names()) == 18
         assert len(trial.tensor_names(collection=CollectionKeys.GRADIENTS)) == 4
-        assert len(trial.tensor_names(collection=CollectionKeys.OPTIMIZER_VARIABLES)) == 5
+    assert len(trial.tensor_names(collection=CollectionKeys.OPTIMIZER_VARIABLES)) == 5
+    assert len(trial.tensor_names(collection="custom_optimizer_variables")) == 5
     assert len(trial.tensor_names(collection=CollectionKeys.BIASES)) == 2
     assert len(trial.tensor_names(collection=CollectionKeys.WEIGHTS)) == 2
     assert len(trial.tensor_names(collection=CollectionKeys.LOSSES)) == 1
@@ -566,6 +580,24 @@ def test_include_collections(out_dir, tf_eager_mode):
     )
 
 
+@pytest.mark.slow
+def test_include_only_custom_collection(out_dir, tf_eager_mode):
+    include_collections = ["custom_optimizer_variables"]
+    save_config = SaveConfig(save_interval=3)
+    hook = smd.KerasHook(
+        out_dir,
+        save_config=save_config,
+        include_collections=include_collections,
+        reduction_config=ReductionConfig(norms=ALLOWED_NORMS, reductions=ALLOWED_REDUCTIONS),
+    )
+    hook.get_collection("custom_optimizer_variables").include("Adam")
+    helper_keras_fit(out_dir, hook=hook, steps=["train", "eval", "predict"], eager=tf_eager_mode)
+
+    trial = smd.create_trial(path=out_dir)
+    assert len(trial.tensor_names()) == (8 if is_tf_2_2() and tf_eager_mode else 9)
+    assert len(trial.tensor_names(collection="custom_optimizer_variables")) == 5
+
+
 @pytest.mark.slow
 def test_hook_from_json(out_dir, tf_eager_mode, monkeypatch):
     monkeypatch.setenv(

tests/tensorflow2/test_keras_mirrored.py

@@ -158,9 +158,9 @@ def exhaustive_check(trial_dir, include_workers="one", eager=True):
     if include_workers == "all":
         assert len(tr.workers()) == strategy.num_replicas_in_sync
         if eager:
-            assert len(tr.tensor_names()) == (6 + 1 + 2 if is_tf_2_2() else 6 + 1 + 3)
-            # 6 weights, 1 loss, 3 metrics for Tf 2.1
-            # 6 weights, 1 loss, 2 metrics for Tf 2.2
+            assert len(tr.tensor_names()) == (6 + 1 + 2 + 5 if is_tf_2_2() else 6 + 1 + 3 + 5)
+            # 6 weights, 1 loss, 3 metrics, 5 optimizer variables for Tf 2.1
+            # 6 weights, 1 loss, 2 metrics, 5 optimizer variables for Tf 2.2
         else:
             assert len(tr.tensor_names()) == (6 + 6 + 1 + 3 + strategy.num_replicas_in_sync * 3 + 5)
     else:
@@ -245,8 +245,8 @@ def test_save_all(out_dir, tf_eager_mode):
     tr = create_trial_fast_refresh(out_dir)
     print(tr.tensor_names())
     if tf_eager_mode:
-        assert len(tr.tensor_names()) == (6 + 2 + 1 if is_tf_2_2() else 6 + 3 + 1)
-        # weights, metrics, losses
+        assert len(tr.tensor_names()) == (6 + 2 + 1 + 5 if is_tf_2_2() else 6 + 3 + 1 + 5)
+        # weights, metrics, losses, optimizer variables
     else:
         assert (
             len(tr.tensor_names())