From d34e39a6736b1374cd65c91094546d64bdefd7e5 Mon Sep 17 00:00:00 2001 From: NihalHarish Date: Thu, 15 Oct 2020 00:56:05 -0700 Subject: [PATCH 01/26] save nested layers --- smdebug/tensorflow/keras.py | 2 +- tests/tensorflow2/test_nested_layers.py | 47 +++++++++++++++++++++++++ 2 files changed, 48 insertions(+), 1 deletion(-) create mode 100644 tests/tensorflow2/test_nested_layers.py diff --git a/smdebug/tensorflow/keras.py b/smdebug/tensorflow/keras.py index 30ba008c9..bf4fb3cfe 100644 --- a/smdebug/tensorflow/keras.py +++ b/smdebug/tensorflow/keras.py @@ -708,7 +708,7 @@ def on_predict_begin(self, logs=None): def _wrap_model_with_input_output_saver(self): if self.has_registered_model: return - for layer in self.model.layers: + for layer in self.model._flatten_layers(include_self=False, recursive=True): layer._hooks = [] layer.call = get_layer_call_fn(layer) layer.register_hook = lambda hook: layer._hooks.append(hook) diff --git a/tests/tensorflow2/test_nested_layers.py b/tests/tensorflow2/test_nested_layers.py new file mode 100644 index 000000000..d02a78bac --- /dev/null +++ b/tests/tensorflow2/test_nested_layers.py @@ -0,0 +1,47 @@ +# Third Party +import numpy as np +from tensorflow.keras.layers import Concatenate, Dense, Layer +from tensorflow.python.keras.models import Model + +# First Party +import smdebug.tensorflow as smd +from smdebug.trials import create_trial + + +class CustomLayer(Layer): + def __init__(self): + super(CustomLayer, self).__init__() + self.con = Concatenate() + self.dense = Dense(10, activation="relu") + + def call(self, x): + x = self.con([x, x]) + return self.dense(x) + + +class CustomModel(Model): + def __init__(self): + super(CustomModel, self).__init__() + self.custom_layer = CustomLayer() + + def call(self, x): + return self.custom_layer(x) + + +def test_if_nested_layers_are_recorded(out_dir): + model = CustomModel() + hook = smd.KerasHook( + out_dir, + save_all=True, + save_config=smd.SaveConfig(save_steps=[0], save_interval=1), + reduction_config=smd.ReductionConfig(save_shape=True, save_raw_tensor=True), + ) + + hook.register_model(model) + x_train = np.random.random((1000, 20)) + y_train = np.random.random((1000, 1)) + model.compile(optimizer="Adam", loss="mse", run_eagerly=True) + model.fit(x_train, y_train, epochs=1, steps_per_epoch=1, callbacks=[hook]) + trial = create_trial(path=out_dir) + layer_names = trial.tensor_names(collection=smd.CollectionKeys.LAYERS) + assert len(layer_names) == 6 From 1d1092a8e7ba6087d8c6e9d2e5bd34c5572da8d9 Mon Sep 17 00:00:00 2001 From: NihalHarish Date: Thu, 15 Oct 2020 15:21:13 -0700 Subject: [PATCH 02/26] comment --- tests/tensorflow2/test_nested_layers.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/tensorflow2/test_nested_layers.py b/tests/tensorflow2/test_nested_layers.py index d02a78bac..44264019e 100644 --- a/tests/tensorflow2/test_nested_layers.py +++ b/tests/tensorflow2/test_nested_layers.py @@ -44,4 +44,6 @@ def test_if_nested_layers_are_recorded(out_dir): model.fit(x_train, y_train, epochs=1, steps_per_epoch=1, callbacks=[hook]) trial = create_trial(path=out_dir) layer_names = trial.tensor_names(collection=smd.CollectionKeys.LAYERS) - assert len(layer_names) == 6 + assert ( + len(layer_names) == 6 + ) # ['custom_layer/inputs', 'custom_layer/outputs', 'con/inputs', 'con/outputs', 'dense/inputs', 'dense/outputs'] From 0edcfa442d24845491211ff3b3efb8d46d193f64 Mon Sep 17 00:00:00 2001 From: NihalHarish Date: Fri, 16 Oct 2020 19:13:32 -0700 Subject: [PATCH 03/26] tmp --- smdebug/tensorflow/keras.py | 7 +++++-- smdebug/tensorflow/utils.py | 12 ++++++------ 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/smdebug/tensorflow/keras.py b/smdebug/tensorflow/keras.py index bf4fb3cfe..db2753fd1 100644 --- a/smdebug/tensorflow/keras.py +++ b/smdebug/tensorflow/keras.py @@ -549,8 +549,11 @@ def _save_layer_input_and_outputs(self): return for layer_name in self.saved_layers: # Save Input - tensor = self.saved_layers[layer_name].layer_input - export_name = get_export_name_for_keras(layer_name, tensor_type="input", tensor=tensor) + layer_inputs = self.saved_layers[layer_name].layer_input + for idx, tensor in layer_inputs: + export_name = get_export_name_for_keras( + layer_name, tensor_type="input", tensor=tensor + ) input_collection = ( {self.get_collection(CollectionKeys.LAYERS)} if self._is_collection_being_saved_for_step(CollectionKeys.LAYERS) diff --git a/smdebug/tensorflow/utils.py b/smdebug/tensorflow/utils.py index 863a5f862..13359429f 100644 --- a/smdebug/tensorflow/utils.py +++ b/smdebug/tensorflow/utils.py @@ -278,11 +278,11 @@ def is_keras_optimizer(obj): return False -def get_export_name_for_keras(layer, tensor_type, tensor=None): +def get_export_name_for_keras(layer, tensor_type, tensor=None, idx=None): if tensor_type in ["input", "output", "weight"]: if isinstance(layer, str): # Tensor.name is meaningless when eager execution is enabled. - return f"{layer}/{tensor_type}s" + return f"{layer}/{tensor_type}_{idx}" else: return f"{layer.name}/{tensor_type}s/{tensor.name}" else: @@ -326,12 +326,12 @@ def register_hook(self, hook: Callable[[tf.Tensor, tf.Tensor], Optional[tf.Tenso class InputOutputSaver: def __init__(self): - self.layer_input = None - self.layer_output = None + self.layer_input = [] + self.layer_output = [] def __call__(self, inputs, *args, **kwargs) -> None: - self.layer_input = kwargs["layer_input"] - self.layer_output = kwargs["layer_output"] + self.layer_input.append(kwargs["layer_input"]) + self.layer_output.append(kwargs["layer_output"]) def get_layer_call_fn(layer: tf.keras.layers.Layer) -> Callable[[tf.Tensor], tf.Tensor]: From 1aa1cfa258e054bed524de551854b25192242b87 Mon Sep 17 00:00:00 2001 From: NihalHarish Date: Sat, 17 Oct 2020 15:55:26 -0700 Subject: [PATCH 04/26] loop --- smdebug/tensorflow/keras.py | 52 ++++++++++++++++++------------------- 1 file changed, 26 insertions(+), 26 deletions(-) diff --git a/smdebug/tensorflow/keras.py b/smdebug/tensorflow/keras.py index db2753fd1..986cbef23 100644 --- a/smdebug/tensorflow/keras.py +++ b/smdebug/tensorflow/keras.py @@ -547,39 +547,39 @@ def _save_metrics(self, batch, logs, force_save=False): def _save_layer_input_and_outputs(self): if is_tf_version_2x() is False: return + input_collection = ( + {self.get_collection(CollectionKeys.LAYERS)} + if self._is_collection_being_saved_for_step(CollectionKeys.LAYERS) + else set() + ) for layer_name in self.saved_layers: # Save Input layer_inputs = self.saved_layers[layer_name].layer_input - for idx, tensor in layer_inputs: + for idx, tensor in enumerate(layer_inputs): export_name = get_export_name_for_keras( - layer_name, tensor_type="input", tensor=tensor + layer_name, tensor_type="input", tensor=tensor, idx=idx ) - input_collection = ( - {self.get_collection(CollectionKeys.LAYERS)} - if self._is_collection_being_saved_for_step(CollectionKeys.LAYERS) - else set() - ) - t = tensor[0] if isinstance(tensor, list) and len(tensor) else tensor - if hasattr(t, "numpy") is False: - self.logger.warning("cannot save layer values during forward pass with tf.function") - continue - else: - self._save_tensor_to_file(export_name, tensor, input_collection) + if hasattr(tensor, "numpy") is False: + self.logger.warning( + "cannot save layer values during forward pass with tf.function" + ) + continue + else: + self._save_tensor_to_file(export_name, tensor, input_collection) # Save Output - tensor = self.saved_layers[layer_name].layer_output - export_name = get_export_name_for_keras(layer_name, tensor_type="output", tensor=tensor) - self._is_collection_being_saved_for_step(CollectionKeys.LAYERS) - output_collection = ( - {self.get_collection(CollectionKeys.LAYERS)} - if self._is_collection_being_saved_for_step(CollectionKeys.LAYERS) - else set() - ) - t = tensor[0] if isinstance(tensor, list) and len(tensor) else tensor - if hasattr(t, "numpy") is False: - self.logger.warning("cannot save layer values during forward pass with tf.function") - else: - self._save_tensor_to_file(export_name, tensor, output_collection) + layer_outputs = self.saved_layers[layer_name].layer_output + for idx, tensor in enumerate(layer_outputs): + export_name = get_export_name_for_keras( + layer_name, tensor_type="output", tensor=tensor, idx=idx + ) + if hasattr(tensor, "numpy") is False: + self.logger.warning( + "cannot save layer values during forward pass with tf.function" + ) + continue + else: + self._save_tensor_to_file(export_name, tensor, input_collection) def _save_tensors_post_step(self, batch, logs): # some tensors available as value from within hook are saved here From d33823e95f3a9870d94bba8b186bbcdde39b3fa1 Mon Sep 17 00:00:00 2001 From: NihalHarish Date: Thu, 15 Oct 2020 00:56:05 -0700 Subject: [PATCH 05/26] save nested layers --- smdebug/tensorflow/keras.py | 2 +- tests/tensorflow2/test_nested_layers.py | 47 +++++++++++++++++++++++++ 2 files changed, 48 insertions(+), 1 deletion(-) create mode 100644 tests/tensorflow2/test_nested_layers.py diff --git a/smdebug/tensorflow/keras.py b/smdebug/tensorflow/keras.py index 30ba008c9..bf4fb3cfe 100644 --- a/smdebug/tensorflow/keras.py +++ b/smdebug/tensorflow/keras.py @@ -708,7 +708,7 @@ def on_predict_begin(self, logs=None): def _wrap_model_with_input_output_saver(self): if self.has_registered_model: return - for layer in self.model.layers: + for layer in self.model._flatten_layers(include_self=False, recursive=True): layer._hooks = [] layer.call = get_layer_call_fn(layer) layer.register_hook = lambda hook: layer._hooks.append(hook) diff --git a/tests/tensorflow2/test_nested_layers.py b/tests/tensorflow2/test_nested_layers.py new file mode 100644 index 000000000..d02a78bac --- /dev/null +++ b/tests/tensorflow2/test_nested_layers.py @@ -0,0 +1,47 @@ +# Third Party +import numpy as np +from tensorflow.keras.layers import Concatenate, Dense, Layer +from tensorflow.python.keras.models import Model + +# First Party +import smdebug.tensorflow as smd +from smdebug.trials import create_trial + + +class CustomLayer(Layer): + def __init__(self): + super(CustomLayer, self).__init__() + self.con = Concatenate() + self.dense = Dense(10, activation="relu") + + def call(self, x): + x = self.con([x, x]) + return self.dense(x) + + +class CustomModel(Model): + def __init__(self): + super(CustomModel, self).__init__() + self.custom_layer = CustomLayer() + + def call(self, x): + return self.custom_layer(x) + + +def test_if_nested_layers_are_recorded(out_dir): + model = CustomModel() + hook = smd.KerasHook( + out_dir, + save_all=True, + save_config=smd.SaveConfig(save_steps=[0], save_interval=1), + reduction_config=smd.ReductionConfig(save_shape=True, save_raw_tensor=True), + ) + + hook.register_model(model) + x_train = np.random.random((1000, 20)) + y_train = np.random.random((1000, 1)) + model.compile(optimizer="Adam", loss="mse", run_eagerly=True) + model.fit(x_train, y_train, epochs=1, steps_per_epoch=1, callbacks=[hook]) + trial = create_trial(path=out_dir) + layer_names = trial.tensor_names(collection=smd.CollectionKeys.LAYERS) + assert len(layer_names) == 6 From 8ac2e626760b5c230c463d3ed5fd73e2cb1dd103 Mon Sep 17 00:00:00 2001 From: NihalHarish Date: Thu, 15 Oct 2020 15:21:13 -0700 Subject: [PATCH 06/26] comment --- tests/tensorflow2/test_nested_layers.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/tensorflow2/test_nested_layers.py b/tests/tensorflow2/test_nested_layers.py index d02a78bac..44264019e 100644 --- a/tests/tensorflow2/test_nested_layers.py +++ b/tests/tensorflow2/test_nested_layers.py @@ -44,4 +44,6 @@ def test_if_nested_layers_are_recorded(out_dir): model.fit(x_train, y_train, epochs=1, steps_per_epoch=1, callbacks=[hook]) trial = create_trial(path=out_dir) layer_names = trial.tensor_names(collection=smd.CollectionKeys.LAYERS) - assert len(layer_names) == 6 + assert ( + len(layer_names) == 6 + ) # ['custom_layer/inputs', 'custom_layer/outputs', 'con/inputs', 'con/outputs', 'dense/inputs', 'dense/outputs'] From 9fb8c45a75c56b1742398f285abd4b4bdab94606 Mon Sep 17 00:00:00 2001 From: NihalHarish Date: Fri, 16 Oct 2020 19:13:32 -0700 Subject: [PATCH 07/26] tmp --- smdebug/tensorflow/keras.py | 7 +++++-- smdebug/tensorflow/utils.py | 12 ++++++------ 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/smdebug/tensorflow/keras.py b/smdebug/tensorflow/keras.py index bf4fb3cfe..db2753fd1 100644 --- a/smdebug/tensorflow/keras.py +++ b/smdebug/tensorflow/keras.py @@ -549,8 +549,11 @@ def _save_layer_input_and_outputs(self): return for layer_name in self.saved_layers: # Save Input - tensor = self.saved_layers[layer_name].layer_input - export_name = get_export_name_for_keras(layer_name, tensor_type="input", tensor=tensor) + layer_inputs = self.saved_layers[layer_name].layer_input + for idx, tensor in layer_inputs: + export_name = get_export_name_for_keras( + layer_name, tensor_type="input", tensor=tensor + ) input_collection = ( {self.get_collection(CollectionKeys.LAYERS)} if self._is_collection_being_saved_for_step(CollectionKeys.LAYERS) diff --git a/smdebug/tensorflow/utils.py b/smdebug/tensorflow/utils.py index 863a5f862..13359429f 100644 --- a/smdebug/tensorflow/utils.py +++ b/smdebug/tensorflow/utils.py @@ -278,11 +278,11 @@ def is_keras_optimizer(obj): return False -def get_export_name_for_keras(layer, tensor_type, tensor=None): +def get_export_name_for_keras(layer, tensor_type, tensor=None, idx=None): if tensor_type in ["input", "output", "weight"]: if isinstance(layer, str): # Tensor.name is meaningless when eager execution is enabled. - return f"{layer}/{tensor_type}s" + return f"{layer}/{tensor_type}_{idx}" else: return f"{layer.name}/{tensor_type}s/{tensor.name}" else: @@ -326,12 +326,12 @@ def register_hook(self, hook: Callable[[tf.Tensor, tf.Tensor], Optional[tf.Tenso class InputOutputSaver: def __init__(self): - self.layer_input = None - self.layer_output = None + self.layer_input = [] + self.layer_output = [] def __call__(self, inputs, *args, **kwargs) -> None: - self.layer_input = kwargs["layer_input"] - self.layer_output = kwargs["layer_output"] + self.layer_input.append(kwargs["layer_input"]) + self.layer_output.append(kwargs["layer_output"]) def get_layer_call_fn(layer: tf.keras.layers.Layer) -> Callable[[tf.Tensor], tf.Tensor]: From e34d80d3ee323f4bd7c22b98f4012fa4fd4bdced Mon Sep 17 00:00:00 2001 From: NihalHarish Date: Sat, 17 Oct 2020 15:55:26 -0700 Subject: [PATCH 08/26] loop --- smdebug/tensorflow/keras.py | 52 ++++++++++++++++++------------------- 1 file changed, 26 insertions(+), 26 deletions(-) diff --git a/smdebug/tensorflow/keras.py b/smdebug/tensorflow/keras.py index db2753fd1..986cbef23 100644 --- a/smdebug/tensorflow/keras.py +++ b/smdebug/tensorflow/keras.py @@ -547,39 +547,39 @@ def _save_metrics(self, batch, logs, force_save=False): def _save_layer_input_and_outputs(self): if is_tf_version_2x() is False: return + input_collection = ( + {self.get_collection(CollectionKeys.LAYERS)} + if self._is_collection_being_saved_for_step(CollectionKeys.LAYERS) + else set() + ) for layer_name in self.saved_layers: # Save Input layer_inputs = self.saved_layers[layer_name].layer_input - for idx, tensor in layer_inputs: + for idx, tensor in enumerate(layer_inputs): export_name = get_export_name_for_keras( - layer_name, tensor_type="input", tensor=tensor + layer_name, tensor_type="input", tensor=tensor, idx=idx ) - input_collection = ( - {self.get_collection(CollectionKeys.LAYERS)} - if self._is_collection_being_saved_for_step(CollectionKeys.LAYERS) - else set() - ) - t = tensor[0] if isinstance(tensor, list) and len(tensor) else tensor - if hasattr(t, "numpy") is False: - self.logger.warning("cannot save layer values during forward pass with tf.function") - continue - else: - self._save_tensor_to_file(export_name, tensor, input_collection) + if hasattr(tensor, "numpy") is False: + self.logger.warning( + "cannot save layer values during forward pass with tf.function" + ) + continue + else: + self._save_tensor_to_file(export_name, tensor, input_collection) # Save Output - tensor = self.saved_layers[layer_name].layer_output - export_name = get_export_name_for_keras(layer_name, tensor_type="output", tensor=tensor) - self._is_collection_being_saved_for_step(CollectionKeys.LAYERS) - output_collection = ( - {self.get_collection(CollectionKeys.LAYERS)} - if self._is_collection_being_saved_for_step(CollectionKeys.LAYERS) - else set() - ) - t = tensor[0] if isinstance(tensor, list) and len(tensor) else tensor - if hasattr(t, "numpy") is False: - self.logger.warning("cannot save layer values during forward pass with tf.function") - else: - self._save_tensor_to_file(export_name, tensor, output_collection) + layer_outputs = self.saved_layers[layer_name].layer_output + for idx, tensor in enumerate(layer_outputs): + export_name = get_export_name_for_keras( + layer_name, tensor_type="output", tensor=tensor, idx=idx + ) + if hasattr(tensor, "numpy") is False: + self.logger.warning( + "cannot save layer values during forward pass with tf.function" + ) + continue + else: + self._save_tensor_to_file(export_name, tensor, input_collection) def _save_tensors_post_step(self, batch, logs): # some tensors available as value from within hook are saved here From ddf30a2c4e1e88525f2e9c6caa65ca0771a71627 Mon Sep 17 00:00:00 2001 From: NihalHarish Date: Mon, 19 Oct 2020 00:17:33 -0700 Subject: [PATCH 09/26] update test --- tests/tensorflow2/test_concat_layer.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/tensorflow2/test_concat_layer.py b/tests/tensorflow2/test_concat_layer.py index d7b31bd5b..a6c3237ac 100644 --- a/tests/tensorflow2/test_concat_layer.py +++ b/tests/tensorflow2/test_concat_layer.py @@ -34,5 +34,4 @@ def test_multiple_inputs(out_dir): trial = create_trial(path=out_dir) tnames = sorted(trial.tensor_names(collection=smd.CollectionKeys.LAYERS)) assert "concatenate" in tnames[0] - assert len(trial.tensor(tnames[0]).value(0)) == 2 - assert trial.tensor(tnames[0]).shape(0) == (2, 1000, 20) + assert trial.tensor(tnames[0]).shape(0) == (1000, 40) From 934632a4dec6e244c27d6380bcbd1bc8172e85dd Mon Sep 17 00:00:00 2001 From: NihalHarish Date: Mon, 19 Oct 2020 01:25:27 -0700 Subject: [PATCH 10/26] fix clash --- smdebug/tensorflow/keras.py | 12 ++++++++++++ smdebug/tensorflow/utils.py | 2 +- 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/smdebug/tensorflow/keras.py b/smdebug/tensorflow/keras.py index 986cbef23..1693891ac 100644 --- a/smdebug/tensorflow/keras.py +++ b/smdebug/tensorflow/keras.py @@ -81,6 +81,7 @@ def __init__( # this flag indicated to the train_batch_begin callback # the the step was already incremented in the on_train_begin callback self.step_incremented_in_on_train_begin = False + self.has_unwrapped_model_with_input_output_saver = False def _is_not_supported(self): if self.distribution_strategy is None: @@ -713,11 +714,19 @@ def _wrap_model_with_input_output_saver(self): return for layer in self.model._flatten_layers(include_self=False, recursive=True): layer._hooks = [] + layer._old_call = layer.call layer.call = get_layer_call_fn(layer) layer.register_hook = lambda hook: layer._hooks.append(hook) saver = InputOutputSaver() layer.register_hook(saver) self.saved_layers[layer.name] = saver + self.has_wrapped_model_with_input_output_saver = True + + def _unwrap_model_with_input_output_saver(self): + for layer in self.model._flatten_layers(include_self=False, recursive=True): + layer._hooks = [] + layer.call = layer._old_call + self.has_unwrapped_model_with_input_output_saver = False def _on_any_batch_begin(self, batch, mode, logs=None): if self._is_not_supported(): @@ -851,6 +860,9 @@ def _on_any_batch_end(self, batch, mode, logs=None): self._export_model() self._exported_model[self.mode] = True + if is_tf_version_2x(): + self._unwrap_model_with_input_output_saver() + def on_train_batch_end(self, batch, logs=None): self._on_any_batch_end(batch, ModeKeys.TRAIN, logs=logs) diff --git a/smdebug/tensorflow/utils.py b/smdebug/tensorflow/utils.py index 13359429f..a37580747 100644 --- a/smdebug/tensorflow/utils.py +++ b/smdebug/tensorflow/utils.py @@ -282,7 +282,7 @@ def get_export_name_for_keras(layer, tensor_type, tensor=None, idx=None): if tensor_type in ["input", "output", "weight"]: if isinstance(layer, str): # Tensor.name is meaningless when eager execution is enabled. - return f"{layer}/{tensor_type}_{idx}" + return f"{layer}/{tensor_type}_{idx}" if idx is not None else f"{layer}/{tensor_type}" else: return f"{layer.name}/{tensor_type}s/{tensor.name}" else: From fb20c632d9b7811ed760d5dabb1f38e6d066600d Mon Sep 17 00:00:00 2001 From: NihalHarish Date: Mon, 19 Oct 2020 01:29:45 -0700 Subject: [PATCH 11/26] idx --- smdebug/tensorflow/keras.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/smdebug/tensorflow/keras.py b/smdebug/tensorflow/keras.py index 1693891ac..e6d7ac2ef 100644 --- a/smdebug/tensorflow/keras.py +++ b/smdebug/tensorflow/keras.py @@ -791,6 +791,7 @@ def _save_layer_values(self, logs): step_collections = self._get_collections_to_save_for_step() layer_collection = self.get_collection(CollectionKeys.LAYERS) collections_to_write = {layer_collection} if layer_collection in step_collections else set() + self.layer_name_dict = dict() for layer_name, layer_input, layer_output in logs: # Cast layer_name to str since it can also be of type bytes # when run with mirrored strategy @@ -805,9 +806,12 @@ def _save_layer_values(self, logs): # Layer Inputs are flattened and passed as a list into # the next layer. Unpacking it speeds up the _make_numpy fn. layer_input = layer_input[0] - layer_input_tensor_name = get_export_name_for_keras(str(layer_name), "input") + layer_name = str(layer_name) + idx = self.layer_name_dict.get(layer_name, 0) + self.layer_name_dict[layer_name] = idx + 1 + layer_input_tensor_name = get_export_name_for_keras(layer_name, "input", idx=idx) self._save_tensor_to_file(layer_input_tensor_name, layer_input, collections_to_write) - layer_output_tensor_name = get_export_name_for_keras(str(layer_name), "output") + layer_output_tensor_name = get_export_name_for_keras(layer_name, "output", idx=idx) self._save_tensor_to_file(layer_output_tensor_name, layer_output, collections_to_write) def _write_optimizer_variables(self): From 7f1378d9e95e715f4e01a85b5165d15d3d317d21 Mon Sep 17 00:00:00 2001 From: NihalHarish Date: Mon, 19 Oct 2020 02:19:09 -0700 Subject: [PATCH 12/26] update test regex --- tests/tensorflow2/test_keras.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/tensorflow2/test_keras.py b/tests/tensorflow2/test_keras.py index 7228b7ec9..1a4d3e17d 100644 --- a/tests/tensorflow2/test_keras.py +++ b/tests/tensorflow2/test_keras.py @@ -194,7 +194,7 @@ def test_layer_names_gradient_tape(out_dir): tr = create_trial_fast_refresh(out_dir) tnames = tr.tensor_names(collection=CollectionKeys.LAYERS) - pattern = r"^(flatten|dense|dropout)(_\d+)?\/(inputs|outputs)" + pattern = r"^(flatten|dense|dropout)(_\d+)?\/(inputs|outputs)_\d+" for tname in tnames: assert re.match(pattern=pattern, string=tname) is not None @@ -580,7 +580,7 @@ def test_layer_names(out_dir, tf_eager_mode): tr = create_trial_fast_refresh(out_dir) tnames = tr.tensor_names(collection=CollectionKeys.LAYERS) - pattern = r"^(flatten|dense|dropout)(_\d+)?\/(inputs|outputs)" + pattern = r"^(flatten|dense|dropout)(_\d+)?\/(inputs|outputs)_\d+" for tname in tnames: assert re.match(pattern=pattern, string=tname) is not None From 15b35fa750774363d3632d51349b9b2545dc6957 Mon Sep 17 00:00:00 2001 From: NihalHarish Date: Mon, 19 Oct 2020 02:35:31 -0700 Subject: [PATCH 13/26] update regex --- tests/tensorflow2/test_keras.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/tensorflow2/test_keras.py b/tests/tensorflow2/test_keras.py index 1a4d3e17d..e83fdbf51 100644 --- a/tests/tensorflow2/test_keras.py +++ b/tests/tensorflow2/test_keras.py @@ -194,7 +194,7 @@ def test_layer_names_gradient_tape(out_dir): tr = create_trial_fast_refresh(out_dir) tnames = tr.tensor_names(collection=CollectionKeys.LAYERS) - pattern = r"^(flatten|dense|dropout)(_\d+)?\/(inputs|outputs)_\d+" + pattern = r"^(flatten|dense|dropout)(_\d+)?\/(input|output)_\d+" for tname in tnames: assert re.match(pattern=pattern, string=tname) is not None @@ -580,7 +580,7 @@ def test_layer_names(out_dir, tf_eager_mode): tr = create_trial_fast_refresh(out_dir) tnames = tr.tensor_names(collection=CollectionKeys.LAYERS) - pattern = r"^(flatten|dense|dropout)(_\d+)?\/(inputs|outputs)_\d+" + pattern = r"^(flatten|dense|dropout)(_\d+)?\/(input|output)_\d+" for tname in tnames: assert re.match(pattern=pattern, string=tname) is not None From cd2cbfe8e4dcfdffba83961caba9627b29eae5d6 Mon Sep 17 00:00:00 2001 From: NihalHarish Date: Mon, 19 Oct 2020 03:05:49 -0700 Subject: [PATCH 14/26] rename var --- smdebug/tensorflow/keras.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/smdebug/tensorflow/keras.py b/smdebug/tensorflow/keras.py index e6d7ac2ef..20f222b71 100644 --- a/smdebug/tensorflow/keras.py +++ b/smdebug/tensorflow/keras.py @@ -81,7 +81,7 @@ def __init__( # this flag indicated to the train_batch_begin callback # the the step was already incremented in the on_train_begin callback self.step_incremented_in_on_train_begin = False - self.has_unwrapped_model_with_input_output_saver = False + self.has_wrapped_model_with_input_output_saver = False def _is_not_supported(self): if self.distribution_strategy is None: @@ -710,7 +710,7 @@ def on_predict_begin(self, logs=None): self._on_any_mode_begin(ModeKeys.PREDICT) def _wrap_model_with_input_output_saver(self): - if self.has_registered_model: + if self.has_wrapped_model_with_input_output_saver: return for layer in self.model._flatten_layers(include_self=False, recursive=True): layer._hooks = [] @@ -726,7 +726,7 @@ def _unwrap_model_with_input_output_saver(self): for layer in self.model._flatten_layers(include_self=False, recursive=True): layer._hooks = [] layer.call = layer._old_call - self.has_unwrapped_model_with_input_output_saver = False + self.has_wrapped_model_with_input_output_saver = False def _on_any_batch_begin(self, batch, mode, logs=None): if self._is_not_supported(): From f4061536fd78551aaeb99ce57c083bbf99703614 Mon Sep 17 00:00:00 2001 From: NihalHarish Date: Mon, 19 Oct 2020 10:40:12 +0000 Subject: [PATCH 15/26] check --- smdebug/tensorflow/keras.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/smdebug/tensorflow/keras.py b/smdebug/tensorflow/keras.py index 20f222b71..2b3753e3f 100644 --- a/smdebug/tensorflow/keras.py +++ b/smdebug/tensorflow/keras.py @@ -710,7 +710,7 @@ def on_predict_begin(self, logs=None): self._on_any_mode_begin(ModeKeys.PREDICT) def _wrap_model_with_input_output_saver(self): - if self.has_wrapped_model_with_input_output_saver: + if self.has_wrapped_model_with_input_output_saver or self.model is None: return for layer in self.model._flatten_layers(include_self=False, recursive=True): layer._hooks = [] @@ -996,6 +996,7 @@ def run(*args, **kwargs): # this means sometimes collections will be exported after 1 step self.export_collections() self._exported_collections = True + self._wrap_model_with_input_output_saver() return run @@ -1067,6 +1068,7 @@ def run(*args, **kwargs): return self.last_saved_step = self.step + self._unwrap_model_with_input_output_saver() return run From de07dad7dcb978c2b67716b0b7f2fa66446a5162 Mon Sep 17 00:00:00 2001 From: NihalHarish Date: Mon, 19 Oct 2020 04:36:04 -0700 Subject: [PATCH 16/26] rename var --- smdebug/tensorflow/keras.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/smdebug/tensorflow/keras.py b/smdebug/tensorflow/keras.py index 2b3753e3f..b5d8d0274 100644 --- a/smdebug/tensorflow/keras.py +++ b/smdebug/tensorflow/keras.py @@ -548,7 +548,7 @@ def _save_metrics(self, batch, logs, force_save=False): def _save_layer_input_and_outputs(self): if is_tf_version_2x() is False: return - input_collection = ( + layer_collection = ( {self.get_collection(CollectionKeys.LAYERS)} if self._is_collection_being_saved_for_step(CollectionKeys.LAYERS) else set() @@ -566,7 +566,7 @@ def _save_layer_input_and_outputs(self): ) continue else: - self._save_tensor_to_file(export_name, tensor, input_collection) + self._save_tensor_to_file(export_name, tensor, layer_collection) # Save Output layer_outputs = self.saved_layers[layer_name].layer_output @@ -580,7 +580,7 @@ def _save_layer_input_and_outputs(self): ) continue else: - self._save_tensor_to_file(export_name, tensor, input_collection) + self._save_tensor_to_file(export_name, tensor, layer_collection) def _save_tensors_post_step(self, batch, logs): # some tensors available as value from within hook are saved here From d698ca5d8021f18a91affc0ab0a77f4740412e86 Mon Sep 17 00:00:00 2001 From: NihalHarish Date: Mon, 19 Oct 2020 05:26:36 -0700 Subject: [PATCH 17/26] self --- smdebug/rules/action/stop_training_action.py | 2 +- smdebug/tensorflow/keras.py | 51 +++++++++++++------- smdebug/tensorflow/utils.py | 4 +- tests/tensorflow2/test_concat_layer.py | 2 +- 4 files changed, 38 insertions(+), 21 deletions(-) diff --git a/smdebug/rules/action/stop_training_action.py b/smdebug/rules/action/stop_training_action.py index c60ac60fb..230096637 100644 --- a/smdebug/rules/action/stop_training_action.py +++ b/smdebug/rules/action/stop_training_action.py @@ -79,7 +79,7 @@ def _get_sm_tj_jobs_with_prefix(self): ) break - return list(found_job_dict.keys()) + return found_job_dict.keys() def _stop_training_job(self): if len(self._found_jobs) != 1: diff --git a/smdebug/tensorflow/keras.py b/smdebug/tensorflow/keras.py index b5d8d0274..c1c14cce6 100644 --- a/smdebug/tensorflow/keras.py +++ b/smdebug/tensorflow/keras.py @@ -556,31 +556,48 @@ def _save_layer_input_and_outputs(self): for layer_name in self.saved_layers: # Save Input layer_inputs = self.saved_layers[layer_name].layer_input - for idx, tensor in enumerate(layer_inputs): - export_name = get_export_name_for_keras( - layer_name, tensor_type="input", tensor=tensor, idx=idx - ) - if hasattr(tensor, "numpy") is False: + for layer_idx, tensor in enumerate(layer_inputs): + if isinstance(tensor, list): + tensor_list = tensor + else: + tensor_list = [tensor] + if hasattr(tensor_list[0], "numpy") is False: self.logger.warning( "cannot save layer values during forward pass with tf.function" ) continue else: - self._save_tensor_to_file(export_name, tensor, layer_collection) - + for t_idx, t in enumerate(tensor_list): + export_name = get_export_name_for_keras( + layer_name, + tensor_type="input", + tensor=tensor, + layer_idx=layer_idx, + tensor_idx=t_idx, + ) + self._save_tensor_to_file(export_name, t, layer_collection) # Save Output - layer_outputs = self.saved_layers[layer_name].layer_output - for idx, tensor in enumerate(layer_outputs): - export_name = get_export_name_for_keras( - layer_name, tensor_type="output", tensor=tensor, idx=idx - ) - if hasattr(tensor, "numpy") is False: + layer_inputs = self.saved_layers[layer_name].layer_input + for layer_idx, tensor in enumerate(layer_inputs): + if isinstance(tensor, list): + tensor_list = tensor + else: + tensor_list = [tensor] + if hasattr(tensor_list[0], "numpy") is False: self.logger.warning( "cannot save layer values during forward pass with tf.function" ) continue else: - self._save_tensor_to_file(export_name, tensor, layer_collection) + for t_idx, t in enumerate(tensor_list): + export_name = get_export_name_for_keras( + layer_name, + tensor_type="output", + tensor=tensor, + layer_idx=layer_idx, + tensor_idx=t_idx, + ) + self._save_tensor_to_file(export_name, t, layer_collection) def _save_tensors_post_step(self, batch, logs): # some tensors available as value from within hook are saved here @@ -791,7 +808,7 @@ def _save_layer_values(self, logs): step_collections = self._get_collections_to_save_for_step() layer_collection = self.get_collection(CollectionKeys.LAYERS) collections_to_write = {layer_collection} if layer_collection in step_collections else set() - self.layer_name_dict = dict() + layer_name_dict = dict() for layer_name, layer_input, layer_output in logs: # Cast layer_name to str since it can also be of type bytes # when run with mirrored strategy @@ -807,8 +824,8 @@ def _save_layer_values(self, logs): # the next layer. Unpacking it speeds up the _make_numpy fn. layer_input = layer_input[0] layer_name = str(layer_name) - idx = self.layer_name_dict.get(layer_name, 0) - self.layer_name_dict[layer_name] = idx + 1 + idx = layer_name_dict.get(layer_name, 0) + layer_name_dict[layer_name] = idx + 1 layer_input_tensor_name = get_export_name_for_keras(layer_name, "input", idx=idx) self._save_tensor_to_file(layer_input_tensor_name, layer_input, collections_to_write) layer_output_tensor_name = get_export_name_for_keras(layer_name, "output", idx=idx) diff --git a/smdebug/tensorflow/utils.py b/smdebug/tensorflow/utils.py index a37580747..d89f8cd58 100644 --- a/smdebug/tensorflow/utils.py +++ b/smdebug/tensorflow/utils.py @@ -278,11 +278,11 @@ def is_keras_optimizer(obj): return False -def get_export_name_for_keras(layer, tensor_type, tensor=None, idx=None): +def get_export_name_for_keras(layer, tensor_type, tensor=None, layer_idx=None, tensor_idx=None): if tensor_type in ["input", "output", "weight"]: if isinstance(layer, str): # Tensor.name is meaningless when eager execution is enabled. - return f"{layer}/{tensor_type}_{idx}" if idx is not None else f"{layer}/{tensor_type}" + return f"{layer}_{layer_idx}/{tensor_type}_{tensor_idx}" else: return f"{layer.name}/{tensor_type}s/{tensor.name}" else: diff --git a/tests/tensorflow2/test_concat_layer.py b/tests/tensorflow2/test_concat_layer.py index a6c3237ac..73a6ab4bc 100644 --- a/tests/tensorflow2/test_concat_layer.py +++ b/tests/tensorflow2/test_concat_layer.py @@ -34,4 +34,4 @@ def test_multiple_inputs(out_dir): trial = create_trial(path=out_dir) tnames = sorted(trial.tensor_names(collection=smd.CollectionKeys.LAYERS)) assert "concatenate" in tnames[0] - assert trial.tensor(tnames[0]).shape(0) == (1000, 40) + assert trial.tensor(tnames[0]).shape(0) == (1000, 20) From 4f1368fe9dc75e8c389dc6d84723989b8dbac02e Mon Sep 17 00:00:00 2001 From: NihalHarish Date: Mon, 19 Oct 2020 05:38:02 -0700 Subject: [PATCH 18/26] idx --- smdebug/tensorflow/keras.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/smdebug/tensorflow/keras.py b/smdebug/tensorflow/keras.py index c1c14cce6..37426964e 100644 --- a/smdebug/tensorflow/keras.py +++ b/smdebug/tensorflow/keras.py @@ -826,9 +826,13 @@ def _save_layer_values(self, logs): layer_name = str(layer_name) idx = layer_name_dict.get(layer_name, 0) layer_name_dict[layer_name] = idx + 1 - layer_input_tensor_name = get_export_name_for_keras(layer_name, "input", idx=idx) + layer_input_tensor_name = get_export_name_for_keras( + layer_name, "input", layer_idx=idx, tensor_idx=idx + ) self._save_tensor_to_file(layer_input_tensor_name, layer_input, collections_to_write) - layer_output_tensor_name = get_export_name_for_keras(layer_name, "output", idx=idx) + layer_output_tensor_name = get_export_name_for_keras( + layer_name, "output", layer_idx=idx, tensor_idx=idx + ) self._save_tensor_to_file(layer_output_tensor_name, layer_output, collections_to_write) def _write_optimizer_variables(self): From 977e9d56d76b10f7645ddd0fb077506447259ef4 Mon Sep 17 00:00:00 2001 From: NihalHarish Date: Mon, 19 Oct 2020 05:53:41 -0700 Subject: [PATCH 19/26] update tests --- tests/tensorflow2/test_keras.py | 4 ++-- tests/tensorflow2/test_nested_layers.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/tensorflow2/test_keras.py b/tests/tensorflow2/test_keras.py index e83fdbf51..051745099 100644 --- a/tests/tensorflow2/test_keras.py +++ b/tests/tensorflow2/test_keras.py @@ -194,7 +194,7 @@ def test_layer_names_gradient_tape(out_dir): tr = create_trial_fast_refresh(out_dir) tnames = tr.tensor_names(collection=CollectionKeys.LAYERS) - pattern = r"^(flatten|dense|dropout)(_\d+)?\/(input|output)_\d+" + pattern = r"^(flatten|dense|dropout)(_\d+)+?\/(input|output)_\d+" for tname in tnames: assert re.match(pattern=pattern, string=tname) is not None @@ -580,7 +580,7 @@ def test_layer_names(out_dir, tf_eager_mode): tr = create_trial_fast_refresh(out_dir) tnames = tr.tensor_names(collection=CollectionKeys.LAYERS) - pattern = r"^(flatten|dense|dropout)(_\d+)?\/(input|output)_\d+" + pattern = r"^(flatten|dense|dropout)(_\d+)+?\/(input|output)_\d+" for tname in tnames: assert re.match(pattern=pattern, string=tname) is not None diff --git a/tests/tensorflow2/test_nested_layers.py b/tests/tensorflow2/test_nested_layers.py index 44264019e..b50f7d459 100644 --- a/tests/tensorflow2/test_nested_layers.py +++ b/tests/tensorflow2/test_nested_layers.py @@ -45,5 +45,5 @@ def test_if_nested_layers_are_recorded(out_dir): trial = create_trial(path=out_dir) layer_names = trial.tensor_names(collection=smd.CollectionKeys.LAYERS) assert ( - len(layer_names) == 6 + len(layer_names) == 8 ) # ['custom_layer/inputs', 'custom_layer/outputs', 'con/inputs', 'con/outputs', 'dense/inputs', 'dense/outputs'] From 824e7313a63c4aac3c7cf14141ab15c4bf628364 Mon Sep 17 00:00:00 2001 From: NihalHarish Date: Mon, 19 Oct 2020 23:15:44 -0700 Subject: [PATCH 20/26] layer reuse test --- .../test_model_that_reuses_layers.py | 42 +++++++++++++++++++ 1 file changed, 42 insertions(+) create mode 100644 tests/tensorflow2/test_model_that_reuses_layers.py diff --git a/tests/tensorflow2/test_model_that_reuses_layers.py b/tests/tensorflow2/test_model_that_reuses_layers.py new file mode 100644 index 000000000..4575170c8 --- /dev/null +++ b/tests/tensorflow2/test_model_that_reuses_layers.py @@ -0,0 +1,42 @@ +# Third Party +import numpy as np +from tensorflow.keras.layers import Dense +from tensorflow.python.keras.models import Model + +# First Party +import smdebug.tensorflow as smd +from smdebug.trials import create_trial + + +class CustomModel(Model): + def __init__(self): + super(CustomModel, self).__init__() + self.dense = Dense(10, activation="relu") + + def call(self, x): + x = self.dense(x) + x = self.dense(x) + return self.dense(x) + + +def test_layer_reusability(out_dir): + model = CustomModel() + hook = smd.KerasHook( + out_dir, + save_all=True, + save_config=smd.SaveConfig(save_steps=[0], save_interval=1), + reduction_config=smd.ReductionConfig(save_shape=True, save_raw_tensor=True), + ) + + hook.register_model(model) + x_train = np.random.random((1000, 10)) + y_train = np.random.random((1000, 1)) + model.compile(optimizer="Adam", loss="mse", run_eagerly=True) + model.fit(x_train, y_train, epochs=1, steps_per_epoch=1, callbacks=[hook]) + + trial = create_trial(path=out_dir, name="training_run") + tensor_names = trial.tensor_names(collection=smd.CollectionKeys.LAYERS) + assert len(tensor_names) == 6 + for name in tensor_names: + shape = trial.tensor(name).shape(step_num=0) + assert shape == (1000, 10) From 6540252758451e05b1e5062e8b1b746d6c87b659 Mon Sep 17 00:00:00 2001 From: NihalHarish Date: Mon, 19 Oct 2020 23:21:57 -0700 Subject: [PATCH 21/26] comments --- tests/tensorflow2/test_model_that_reuses_layers.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tests/tensorflow2/test_model_that_reuses_layers.py b/tests/tensorflow2/test_model_that_reuses_layers.py index 4575170c8..539e729a0 100644 --- a/tests/tensorflow2/test_model_that_reuses_layers.py +++ b/tests/tensorflow2/test_model_that_reuses_layers.py @@ -36,6 +36,16 @@ def test_layer_reusability(out_dir): trial = create_trial(path=out_dir, name="training_run") tensor_names = trial.tensor_names(collection=smd.CollectionKeys.LAYERS) + """ + [ + 'dense_0/input_0', + 'dense_0/output_0', + 'dense_1/input_0', + 'dense_1/output_0', + 'dense_2/input_0', + 'dense_2/output_0' + ] + """ assert len(tensor_names) == 6 for name in tensor_names: shape = trial.tensor(name).shape(step_num=0) From 613ecf676da47f00013216ffeec90243fe60b593 Mon Sep 17 00:00:00 2001 From: NihalHarish Date: Tue, 20 Oct 2020 00:01:38 -0700 Subject: [PATCH 22/26] comments --- smdebug/tensorflow/keras.py | 6 ++++-- tests/tensorflow2/test_concat_layer.py | 9 +++++++-- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/smdebug/tensorflow/keras.py b/smdebug/tensorflow/keras.py index 37426964e..249ee70f3 100644 --- a/smdebug/tensorflow/keras.py +++ b/smdebug/tensorflow/keras.py @@ -577,8 +577,8 @@ def _save_layer_input_and_outputs(self): ) self._save_tensor_to_file(export_name, t, layer_collection) # Save Output - layer_inputs = self.saved_layers[layer_name].layer_input - for layer_idx, tensor in enumerate(layer_inputs): + layer_outputs = self.saved_layers[layer_name].layer_output + for layer_idx, tensor in enumerate(layer_outputs): if isinstance(tensor, list): tensor_list = tensor else: @@ -728,6 +728,8 @@ def on_predict_begin(self, logs=None): def _wrap_model_with_input_output_saver(self): if self.has_wrapped_model_with_input_output_saver or self.model is None: + # do not proceed if the model has already been wrapped + # or the model has not been registered with smdebug yet return for layer in self.model._flatten_layers(include_self=False, recursive=True): layer._hooks = [] diff --git a/tests/tensorflow2/test_concat_layer.py b/tests/tensorflow2/test_concat_layer.py index 73a6ab4bc..b68e3850d 100644 --- a/tests/tensorflow2/test_concat_layer.py +++ b/tests/tensorflow2/test_concat_layer.py @@ -32,6 +32,11 @@ def test_multiple_inputs(out_dir): my_model.fit(x_train, y_train, epochs=1, steps_per_epoch=1, callbacks=[hook]) trial = create_trial(path=out_dir) - tnames = sorted(trial.tensor_names(collection=smd.CollectionKeys.LAYERS)) - assert "concatenate" in tnames[0] + tnames = trial.tensor_names(regex="concatenate") + assert len(tnames) == 3 # two inputs + one output + tnames = trial.tensor_names(regex="concatenate.+/input") + assert len(tnames) == 2 # Concatenate Layer receives two inputs assert trial.tensor(tnames[0]).shape(0) == (1000, 20) + tnames = trial.tensor_names(regex="concatenate.+/output") + assert len(tnames) == 1 # Concatenate Layer emits a single output + assert trial.tensor(tnames[0]).shape(0) == (1000, 40) From 5e613c3ea9514f3c8410aef09d34ac5b8344d1ce Mon Sep 17 00:00:00 2001 From: NihalHarish Date: Tue, 20 Oct 2020 00:07:32 -0700 Subject: [PATCH 23/26] revert --- smdebug/rules/action/stop_training_action.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/smdebug/rules/action/stop_training_action.py b/smdebug/rules/action/stop_training_action.py index 230096637..c60ac60fb 100644 --- a/smdebug/rules/action/stop_training_action.py +++ b/smdebug/rules/action/stop_training_action.py @@ -79,7 +79,7 @@ def _get_sm_tj_jobs_with_prefix(self): ) break - return found_job_dict.keys() + return list(found_job_dict.keys()) def _stop_training_job(self): if len(self._found_jobs) != 1: From 3034a12b86cb6e94d2047451220e7cf97d751e24 Mon Sep 17 00:00:00 2001 From: NihalHarish Date: Tue, 20 Oct 2020 00:27:14 -0700 Subject: [PATCH 24/26] comment --- tests/tensorflow2/test_nested_layers.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/tests/tensorflow2/test_nested_layers.py b/tests/tensorflow2/test_nested_layers.py index b50f7d459..79e41d274 100644 --- a/tests/tensorflow2/test_nested_layers.py +++ b/tests/tensorflow2/test_nested_layers.py @@ -44,6 +44,15 @@ def test_if_nested_layers_are_recorded(out_dir): model.fit(x_train, y_train, epochs=1, steps_per_epoch=1, callbacks=[hook]) trial = create_trial(path=out_dir) layer_names = trial.tensor_names(collection=smd.CollectionKeys.LAYERS) - assert ( - len(layer_names) == 8 - ) # ['custom_layer/inputs', 'custom_layer/outputs', 'con/inputs', 'con/outputs', 'dense/inputs', 'dense/outputs'] + """ + [ + 'concatenate_0/input_0', + 'concatenate_0/input_1', + 'concatenate_0/output_0', + 'custom_layer_0/input_0', + 'custom_layer_0/output_0', + 'dense_0/input_0', + 'dense_0/output_0' + ] + """ + assert len(layer_names) == 7 From 277e4cd89f72680d1ebb774e8215d04bea07d445 Mon Sep 17 00:00:00 2001 From: NihalHarish Date: Fri, 23 Oct 2020 19:00:48 +0000 Subject: [PATCH 25/26] check before wrap --- smdebug/tensorflow/keras.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/smdebug/tensorflow/keras.py b/smdebug/tensorflow/keras.py index 249ee70f3..2631326d1 100644 --- a/smdebug/tensorflow/keras.py +++ b/smdebug/tensorflow/keras.py @@ -727,7 +727,7 @@ def on_predict_begin(self, logs=None): self._on_any_mode_begin(ModeKeys.PREDICT) def _wrap_model_with_input_output_saver(self): - if self.has_wrapped_model_with_input_output_saver or self.model is None: + if self.has_wrapped_model_with_input_output_saver or self.model is None or self.has_default_hook_configuration(): # do not proceed if the model has already been wrapped # or the model has not been registered with smdebug yet return @@ -742,6 +742,8 @@ def _wrap_model_with_input_output_saver(self): self.has_wrapped_model_with_input_output_saver = True def _unwrap_model_with_input_output_saver(self): + if self.has_wrapped_model_with_input_output_saver is False: + return for layer in self.model._flatten_layers(include_self=False, recursive=True): layer._hooks = [] layer.call = layer._old_call From 6992be73668fb400efe0be231f6d1e0efba646a0 Mon Sep 17 00:00:00 2001 From: NihalHarish Date: Fri, 23 Oct 2020 12:01:42 -0700 Subject: [PATCH 26/26] lint --- smdebug/tensorflow/keras.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/smdebug/tensorflow/keras.py b/smdebug/tensorflow/keras.py index 2631326d1..a5bf93483 100644 --- a/smdebug/tensorflow/keras.py +++ b/smdebug/tensorflow/keras.py @@ -727,7 +727,11 @@ def on_predict_begin(self, logs=None): self._on_any_mode_begin(ModeKeys.PREDICT) def _wrap_model_with_input_output_saver(self): - if self.has_wrapped_model_with_input_output_saver or self.model is None or self.has_default_hook_configuration(): + if ( + self.has_wrapped_model_with_input_output_saver + or self.model is None + or self.has_default_hook_configuration() + ): # do not proceed if the model has already been wrapped # or the model has not been registered with smdebug yet return