in smdebug/tensorflow/keras.py [0:0]
def _wrap_tape_gradient(self, function):
"""
tape.gradient() is used to compute gradients from loss and model variables.
Using this wrapper to get gradients, loss, weights, and bias values.
"""
@functools.wraps(function)
@error_handling_agent.catch_smdebug_errors(default_return_val=function)
def run(*args, **kwargs):
grads = function(*args, **kwargs)
if self._is_not_supported():
return grads
loss = args[1]
vars = args[2]
if (
(not grads or not vars)
or (not isinstance(grads, list) or not isinstance(vars, list))
or (
not (
(isinstance(vars[0], supported_tf_variables()))
and hasattr(vars[0], "numpy")
)
)
or (not ((isinstance(grads[0], tf.Tensor)) and hasattr(grads[0], "numpy")))
):
return grads
if self._get_collections_to_save_for_step():
for (g, v) in zip(grads, vars):
layer = v.name.split(":")[0]
# Adding a check to make sure gradients are not None.
# gradients may be None if user tries to compute gradients for
# non-training variable when using model.variables instead of
# model.trainable_variables in tape.gradient().
# model.variables includes trainable and non-trainable
# variables.
if g is not None:
self._save_for_tensor(
tensor_name="gradients/" + layer + "Grad",
tensor_value=g,
check_before_write=True,
)
self._save_for_tensor(
tensor_name="weights/" + v.name,
tensor_value=v.value(),
check_before_write=True,
)
self._write_optimizer_variables()
self._save_layer_input_and_outputs()
if not ((isinstance(loss, tf.Tensor)) and hasattr(loss, "numpy")):
return grads
self._add_metric(metric_name="loss", metric_value=loss)
if self._is_collection_being_saved_for_step(CollectionKeys.LOSSES):
self._initialize_writers(only_initialize_if_missing=True)
self._save_for_tensor("loss", loss, check_before_write=False)
return grads
return run