def __init__(self, trackable):
if not isinstance(trackable, tf.__internal__.tracking.Trackable):
raise ValueError(f"{trackable} is not a Trackable object.")
self._trackable = trackable
self._distribute_strategy = tf.distribute.get_strategy()
saveables = tf.__internal__.tracking.saveable_objects_from_trackable(
trackable).values()
# 'Saveables' won't exist when we're passed a legacy TF1 table like
# a StaticHashTable.
if not saveables:
self._num_tensors = 0
self._setter = lambda weights: None
self._getter = lambda: []
elif len(saveables) == 1:
saveable = list(saveables)[0]
if tf1.executing_eagerly_outside_functions():
# If we're in eager mode, we need to defer calling the
# Trackable's saveable() callable until data export time.
# However, it is safe to call the saveable as many times as we
# want, so we will call it now to figure out how many tensors
# this Trackable will produce.
self._saveable = saveable
self._num_tensors = len(self._saveable().specs)
self._setter = lambda weights: self._saveable().restore(
weights, None)
self._getter = lambda: [
spec.tensor for spec in self._saveable().specs
]
else:
# If we're in Graph mode, we need to evaluate the Saveable only
# once and cache the resulting restore graph. Failing to do this
# will result in new assignment ops being added to the graph
# each time set_weights() is called.
self._placeholder_tensors = []
self._saveable = saveable()
self._num_tensors = len(self._saveable.specs)
for spec in self._saveable.specs:
tensor = spec.tensor
self._placeholder_tensors.append(
tf1.placeholder(tensor.dtype, tensor.shape))
self._assign_op = self._saveable.restore(
self._placeholder_tensors, None)
self._setter = self._set_weights_v1
self._getter = lambda: [
spec.tensor for spec in self._saveable.specs
]
else:
raise ValueError(
"Only Trackables with one Saveable are supported. "
f"The Trackable {trackable} has {len(saveables)} Saveables.")
def is_in_tf_function():
"""Returns if inside of a tf.function."""
# Check if running in V1 graph mode.
if not tf1.executing_eagerly_outside_functions():
return False
if not tf.inside_function():
return False
# Check if inside Keras FuncGraph.
if is_in_keras_graph():
return False
# Check for a v1 `wrap_function` FuncGraph.
graph = tf1.get_default_graph()
if getattr(graph, "name",
False) and graph.name.startswith("wrapped_function"):
return False
return True
def apply_gradients(self, grads_and_vars, global_step=None, name=None):
if global_step is None:
global_step = tf.train.get_or_create_global_step()
new_global_step = global_step + 1
assignments = []
for (grad, param) in grads_and_vars:
if grad is None or param is None:
continue
if tf.executing_eagerly() or tf.executing_eagerly_outside_functions():
param_name = param.name
else:
param_name = param.op.name
v = tf.get_variable(
name=f'{param_name}/{self.get_name()}/Momentum',
shape=param.shape.as_list(),
dtype=param.dtype,
trainable=False,
initializer=tf.zeros_initializer())
grad = tf.cast(grad, param.dtype)
if self._use_weight_decay(param_name):
grad += self.weight_decay * param
if callable(self.learning_rate):
base_learning_rate = self.learning_rate()
else:
base_learning_rate = self.learning_rate
if self.classic_momentum:
trust_ratio = 1.0
if self._do_layer_adaptation(param_name):
w_norm = tf.norm(param, ord=2)
g_norm = tf.norm(grad, ord=2)
trust_ratio = tf.where(
tf.greater(w_norm, 0), tf.where(
tf.greater(g_norm, 0), (self.eta * w_norm / g_norm),
1.0),
1.0)
trust_ratio = tf.cast(trust_ratio, param.dtype)
base_learning_rate = tf.cast(base_learning_rate, param.dtype)
scaled_lr = base_learning_rate * trust_ratio
next_v = tf.multiply(
tf.cast(self.momentum, v.dtype), v) + scaled_lr * grad
if self.use_nesterov:
update = tf.multiply(
tf.cast(self.momentum, v.dtype), next_v) + scaled_lr * grad
else:
update = next_v
next_param = param - update
else:
next_v = tf.multiply(
tf.cast(self.momentum, v.dtype), v) + grad
if self.use_nesterov:
update = tf.multiply(
tf.cast(self.momentum, v.dtype), next_v) + grad
else:
update = next_v
trust_ratio = 1.0
if self._do_layer_adaptation(param_name):
w_norm = tf.norm(param, ord=2)
v_norm = tf.norm(update, ord=2)
trust_ratio = tf.where(
tf.greater(w_norm, 0), tf.where(
tf.greater(v_norm, 0), (self.eta * w_norm / v_norm),
1.0),
1.0)
trust_ratio = tf.cast(trust_ratio, param.dtype)
scaled_lr = trust_ratio * base_learning_rate
next_param = param - scaled_lr * update
assignments.extend(
[param.assign(next_param),
v.assign(next_v),
global_step.assign(new_global_step)])
return tf.group(*assignments, name=name)
def check_graph_consistency(tensor=None, method="add_loss", force_raise=False):
"""Checks that tensors passed to `add_*` method match the Keras graph.
When one of the `add_*` method is called inside a V2 conditional branch, the
underlying tensor gets created in a FuncGraph managed by control_flow_v2.
We need to raise clear error messages in such cases.
Args:
tensor: Tensor to check, or `False` if it is known that an error
should be raised.
method: Caller method, one of {'add_metric', 'add_loss', 'add_update'}.
force_raise: If an error should be raised regardless of `tensor`.
Raises:
RuntimeError: In case of an out-of-graph tensor.
"""
if force_raise or (tf1.executing_eagerly_outside_functions() and hasattr(
tensor, "graph") and tensor.graph.is_control_flow_graph):
if method == "activity_regularizer":
bad_example = """
class TestModel(tf.keras.Model):
def __init__(self):
super(TestModel, self).__init__(name='test_model')
self.dense = tf.keras.layers.Dense(2, activity_regularizer='l2')
def call(self, x, training=None):
if training:
return self.dense(x)
else:
return self.dense(x)
"""
correct_example = """
class TestModel(tf.keras.Model):
def __init__(self):
super(TestModel, self).__init__(name='test_model')
self.dense = tf.keras.layers.Dense(2, activity_regularizer='l2')
def call(self, x, training=None):
return self.dense(x)
"""
raise RuntimeError(
"You are using a layer with `activity_regularizer` in a "
f"control flow branch, e.g.:\n{bad_example}\nThis is currently "
"not supported. Please move your call to the layer with "
"`activity_regularizer` out of the control flow branch, "
f"e.g.:\n{correct_example}\nYou can also resolve this by "
"marking your outer model/layer dynamic (eager-only) by "
"passing `dynamic=True` to the layer constructor. Any kind of "
"control flow is supported with dynamic layers. Note that "
"using `dynamic=True` requires you to implement static shape "
"inference in the `compute_output_shape(input_shape)` "
"method.")
if method == "add_metric":
bad_example = """
def call(self, inputs, training=None):
if training:
metric = compute_metric(inputs)
self.add_metric(metric, name='my_metric', aggregation='mean')
return inputs
"""
correct_example = """
def call(self, inputs, training=None):
if training:
metric = compute_metric(inputs)
else:
metric = 0.
self.add_metric(metric, name='my_metric', aggregation='mean')
return inputs
"""
elif method == "add_loss":
bad_example = """
def call(self, inputs, training=None):
if training:
loss = compute_loss(inputs)
self.add_loss(loss)
return inputs
"""
correct_example = """
def call(self, inputs, training=None):
if training:
loss = compute_loss(inputs)
else:
loss = 0.
self.add_loss(loss)
return inputs
"""
else:
bad_example = """
def call(self, inputs, training=None):
if training:
self.add_update(self.w.assign_add(1))
return inputs
"""
correct_example = """
def call(self, inputs, training=None):
if training:
increment = 1
else:
increment = 0
self.add_update(self.w.assign_add(increment))
return inputs
"""
raise RuntimeError(
"You are using the method `{method}` in a control flow branch "
"in your layer, e.g.:\n{bad_example}\n"
"This is not currently supported. "
"Please move your call to {method} out of the control flow branch, "
"e.g.:\n{correct_example}\n"
"You can also resolve this by marking your layer "
"as dynamic (eager-only) by passing "
"`dynamic=True` to the layer constructor. "
"Any kind of control flow is supported with dynamic layers. "
"Note that using `dynamic=True` requires you "
"to implement static shape inference "
"in the `compute_output_shape(input_shape)` method.".format(
method=method,
bad_example=bad_example,
correct_example=correct_example,
))
def _create_keras_history_helper(tensors, processed_ops, created_layers):
"""Helper method for `create_keras_history`.
Args:
tensors: A structure of Tensors for which to create Keras metadata.
processed_ops: Set. TensorFlow operations that have already been wrapped
in `TensorFlowOpLayer` instances.
created_layers: List. The `TensorFlowOpLayer` instances created.
Returns:
Tuple. First element is the updated set of TensorFlow Operations that
have been wrapped in `TensorFlowOpLayer` instances. Second element is
a list of the `TensorFlowOpLayer` instances created.
"""
if tf1.executing_eagerly_outside_functions():
raise ValueError(
"`create_keras_history` should only be called if eager is disabled!"
)
# Import of `base_layer` needed in order to create `TensorFlowOpLayer`.
# Cannot be imported at top because of circular dependencies.
# TODO(omalleyt): Resolve circular dependency.
from keras.engine import base_layer
tensor_list = tf.nest.flatten(tensors)
sparse_ops = []
ragged_tensors = []
for tensor in tensor_list:
if getattr(tensor, "_keras_history", None) is not None:
continue
if isinstance(tensor, (tf.SparseTensor, tf1.SparseTensorValue)):
sparse_ops.append(tensor.op)
continue
if tf_utils.is_ragged(tensor):
# Ragged tensors don't have an op property
ragged_tensors.append(tensor)
continue
op = tensor.op # The Op that created this Tensor.
if op not in processed_ops:
# Recursively set `_keras_history`.
op_inputs = list(op.inputs)
constants = {}
layer_inputs = []
for i, op_input in enumerate(op_inputs):
if uses_keras_history(op_input):
layer_inputs.append(op_input)
else:
# Treat any value not originating from a `keras.Input` as
# a constant. Variables cannot be supported.
ds_with_session = (
tf.distribute.in_cross_replica_context()
and not tf1.executing_eagerly_outside_functions())
using_xla = control_flow_util.GraphOrParentsInXlaContext(
tf1.get_default_graph())
if (ds_with_session or using_xla
or _UNSAFE_GRAPH_OP_LAYER_CREATION):
# In Legacy Graph mode, evaluating here makes Session be
# configured improperly. The downside of this is that
# saving via `get_config` breaks, but SavedModel still
# works.
constants[i] = op_input
else:
with tf.init_scope():
constants[i] = backend.function([], op_input)([])
layer_inputs = unnest_if_single_tensor(layer_inputs)
processed_ops, created_layers = _create_keras_history_helper(
layer_inputs, processed_ops, created_layers)
name = op.name
node_def = op.node_def.SerializeToString()
op_layer = base_layer.TensorFlowOpLayer(node_def,
constants=constants,
name=name)
created_layers.append(op_layer)
op_layer._set_connectivity_metadata(args=(layer_inputs, ),
kwargs={},
outputs=op.outputs)
processed_ops.update([op])
if sparse_ops or ragged_tensors:
lambda_example = """
weights_mult = lambda x: tf.sparse.sparse_dense_matmul(x, weights)
output = tf.keras.layers.Lambda(weights_mult)(input)
"""
raise ValueError(
"Tensorflow ops that generate ragged or sparse tensor "
"outputs are currently not supported by Keras automatic "
"op wrapping. Please wrap these ops in a Lambda layer: "
"\n\n```\n{example}\n```\n"
"Sparse ops encountered: {sparse_ops}\n"
"Ragged tensors encountered: {ragged_tensors}\n".format(
example=lambda_example,
sparse_ops=str(sparse_ops),
ragged_tensors=str(ragged_tensors),
))
return processed_ops, created_layers