def notify_about_new_variables(callback):
"""Calls `callback(var)` for all newly created variables.
Callback should not modify the variable passed in. Use cases that require
variables to be modified should use `variable_creator_scope` directly and sit
within the variable creator stack.
>>> variables = []
>>> with notify_about_variables(variables.append):
... v = tf.Variable(1.0, name='v')
... w = tf.get_variable('w', [])
>>> assert variables == [v, w]
Args:
callback: a callable taking a single argument which is a tf.Variable.
Yields:
`None` - used for contextmanager API.
"""
def _tracking_creator(getter, **kwargs):
v = getter(**kwargs)
callback(v)
return v
with tf.variable_creator_scope(_tracking_creator):
yield
def build_network(self, inputs, phase_train=True, nclass=1001):
try:
from official.recommendation import neumf_model # pylint: disable=g-import-not-at-top
except ImportError as e:
if 'neumf_model' not in e.message:
raise
raise ImportError('To use the experimental NCF model, you must clone the '
'repo https://github.com/tensorflow/models and add '
'tensorflow/models to the PYTHONPATH.')
del nclass
users, items, _ = inputs
params = {
'num_users': _NUM_USERS_20M,
'num_items': _NUM_ITEMS_20M,
'model_layers': (256, 256, 128, 64),
'mf_dim': 64,
'mf_regularization': 0,
'mlp_reg_layers': (0, 0, 0, 0),
'use_tpu': False
}
if self.data_type == tf.float32:
keras_model = neumf_model.construct_model(users, items, params)
logits = keras_model.output
else:
assert self.data_type == tf.float16
tf.keras.backend.set_floatx('float16')
# We cannot rely on the variable_scope's fp16 custom getter here, because
# the NCF model uses keras layers, which ignore variable scopes. So we use
# a variable_creator_scope instead.
with tf.variable_creator_scope(_fp16_variable_creator):
keras_model = neumf_model.construct_model(users, items, params)
logits = tf.cast(keras_model.output, tf.float32)
return model.BuildNetworkResult(logits=logits, extra_info=None)
def __call__(self, image, probe=True, aug_image_key='image'):
# creating local variable which will store copy of CTA log probabilities
with tf.variable_creator_scope(_skip_mirrored_creator):
local_log_prob = tf.Variable(
lambda: tf.ones(self.state_shape, dtype=tf.float32),
trainable=False,
name='cta_log_probs')
self.log_probs.append(local_log_prob)
output_dict = {}
if probe:
probe_op_indices, probe_op_args = self._sample_ops_uniformly()
probe_image = self._apply_ops(image, probe_op_indices,
probe_op_args)
output_dict['probe_op_indices'] = probe_op_indices
output_dict['probe_op_args'] = probe_op_args
output_dict['probe_image'] = probe_image
if aug_image_key is not None:
op_indices, op_args = self._sample_ops(local_log_prob)
aug_image = self._apply_ops(image,
op_indices,
op_args,
prob_to_apply=self.prob_to_apply)
output_dict[aug_image_key] = aug_image
if aug_image_key != 'image':
output_dict['image'] = image
return output_dict
def testOptimizerInsideModelFn(self, distribution, optimizer_fn):
if (not tf.executing_eagerly()
and tf.compat.v1.control_flow_v2_enabled()):
self.skipTest("b/138751864")
created_variables = []
trainable_variables = []
def appending_creator(next_creator, **kwargs):
v = next_creator(**kwargs)
created_variables.append(v.name)
if "trainable" in kwargs and kwargs["trainable"]:
trainable_variables.append(v.name)
return v
# Creator scope needs to be set before it's used inside
# `distribution.scope`.
with tf.variable_creator_scope(
appending_creator), distribution.scope():
optimizer = optimizer_fn()
model_fn, dataset_fn, _ = minimize_loss_example(
optimizer, use_bias=True, use_callable_loss=True)
def step_fn(ctx, inputs):
del ctx # Unused
return distribution.group(
distribution.extended.call_for_each_replica(
model_fn, args=(inputs, )))
iterator = self._get_iterator(distribution, dataset_fn)
def run_step():
return distribution.extended.experimental_run_steps_on_iterator(
step_fn, iterator, iterations=1).run_op
if not tf.executing_eagerly():
with self.cached_session() as sess:
run_step = sess.make_callable(run_step())
self.evaluate(tf.compat.v1.global_variables_initializer())
run_step()
def get_expected_variables(num_parameter_devices):
name = optimizer._name
if isinstance(optimizer, optimizer_v2.OptimizerV2):
variables = VAR_MAP_V2[name]
else:
variables = VAR_MAP_V1[name]
extended_variables = [
v + "/replica_{}".format(replica) for v in variables
for replica in range(1, num_parameter_devices)
]
variables = list(variables) + extended_variables
return set(v + ":0" for v in variables)
self.assertEqual(
get_expected_variables(
len(distribution.extended.parameter_devices)),
set(created_variables))
def test_demo():
with tf.variable_creator_scope('chart1'):
a = tf.Variable(40)
b = tf.Variable(50)
with tf.variable_creator_scope('chart2'):
c = tf.add(a, b)
# init = tf.global_variables_initializer()
print(a)
print(b)
print(c)
with tf.Session() as sess:
# sess.run(init)
a_value, b_value, c_value = sess.run([a, b, c])
print("a_value", a_value)
print("b_value", b_value)
print("c_value", c_value)
return None
def anchor_target_layer(cls_pre, bbox, img_info, scope_name):
""" Get the results of anchor
Args:
cls_pre (float): classifier predict
bbox : bounding boxes
img_info: image information
"""
with tf.variable_creator_scope(scope_name) as scope:
def update(state, y_true, y_pred, sample_weight=None):
del sample_weight # Unused.
def update(metric):
metric.update_state(y_true, y_pred, sample_weight=None)
return metric.variables
with tf.variable_creator_scope(
build_replace_variable_with_parameter_creator(state)):
return tf.nest.map_structure(update, metrics_constructor())
def predict_on_batch(model_weights: ModelWeights,
x: Any,
training: bool = True) -> Any:
with tf.init_scope():
if tf.executing_eagerly():
raise KerasFunctionalModelError(
'tf.keras.Model used as a FunctionalModel is only usable inside a '
'tff.tf_computation decorated callable or a graph context.'
)
# Make a copy of the weights container; can't mutate Python containers
# inside a tf.function.
trainable, non_trainable = (list(w) for w in model_weights)
# Here were intercept variable creation requests during the
# `tf.keras.models.clone_model()` call.
#
# Instead of forwarding the variable request to TF core and getting a
# `tf.Variable` back, we skip that and return only the `tf.Tensor` that
# corresponds to the `tf.Variable` recreation request (avoiding any variable
# creation). This works because TF operations that accept `tf.Variable`
# inputs automatically call `variable.read_value()` and then operate on that
# resulting tensor. We're relying on shortcutting that and providing the
# tensor straight away.
#
# For example, `tf.matmul` doesn't notice its input is `tf.Variable` or
# `tf.Tensor`:
#
# v = tf.Variable([[1], [2], [3]])
# tf.matmul(v, [[4, 5, 6]])
#
# and
#
# v = tf.constant([[1], [2], [3]])
# tf.matmul(v, [[4, 5, 6]])
#
# both result in:
#
# <tf.Tensor: shape=(3, 3), dtype=int32, numpy=
# array([[ 4, 5, 6],
# [ 8, 10, 12],
# [12, 15, 18]], dtype=int32)>
def swap_tensor_parameter_for_variable(_, **kwargs):
if kwargs.get('trainable', True):
return trainable.pop(0)
else:
return non_trainable.pop(0)
with tf.variable_creator_scope(swap_tensor_parameter_for_variable):
if isinstance(keras_model, tf.keras.Model):
variableless_model = tf.keras.models.clone_model(keras_model)
else:
variableless_model = keras_model()
return variableless_model(x, training)
def write_v2_saved_model(tf_function: function.Function, name: str,
saved_model_dir: str) -> function.ConcreteFunction:
"""Writes `tf_function` under attr `name` to `saved_model_dir`."""
module = tf.Module()
resource_tracker = tracking.ResourceTracker()
object_tracker = annotators.ObjectTracker()
created_variables = []
def _variable_creator(next_creator, **kwargs):
var = next_creator(**kwargs)
created_variables.append(var)
return var
# TODO(b/164921571): Handle generic Trackable objects.
# Trace `tf_function` to gather any resources in it using the
# resource_tracker. These are then assigned to `module.resources` and tracked
# before exporting to SavedModel.
with tracking.resource_tracker_scope(resource_tracker), \
annotators.object_tracker_scope(object_tracker), \
tf.variable_creator_scope(_variable_creator):
concrete_fn = tf_function.get_concrete_function()
# Prior to 2020/10/08, saving a tf.function with a concrete function signature
# would ensure that the function was not re-traced in a round-trip to a
# SavedModel. Since this is no longer the case, we save the concrete function
# directly.
if tf.compat.forward_compatible(2020, 10, 8):
pruned_function = optimize_concrete_function(concrete_fn)
module.pruned_variables = pruned_function.variables
setattr(module, name, pruned_function)
else:
setattr(module, name, tf_function)
# Any variables created need to be explicitly tracked.
module.created_variables = created_variables
# Resources need to be explicitly tracked.
module.resources = resource_tracker.resources
module.trackable_objects = object_tracker.trackable_objects
# TODO(b/158011374) - Stop explicitly tracking initializers. Tracking the
# table should be sufficient.
initializers = []
for resource in module.resources:
if isinstance(resource, lookup_ops.InitializableLookupTableBase):
initializers.append(resource._initializer) # pylint: disable=protected-access
module.initializers = initializers
module.assets = [
common_types.Asset(asset_filepath)
for asset_filepath in concrete_fn.graph.get_collection(
tf.compat.v1.GraphKeys.ASSET_FILEPATHS)
]
tf.saved_model.save(module, saved_model_dir)
return concrete_fn
def record_variable_creation_scope():
"""Creates a single use contextmanager for capture variable creation calls."""
variable_list = []
def logging_variable_creator(next_creator, **kwargs):
variable = next_creator(**kwargs)
variable_list.append(variable)
return variable
with contextlib.ExitStack() as stack:
stack.enter_context(
tf.variable_creator_scope(logging_variable_creator))
yield variable_list
def __call__(self,
data: tf.Tensor,
probe: bool = True,
aug_key: str = 'data') -> dict:
"""
When training labeled data, use `probe=True` to update weights
When training unlabeled data, use `probe=False aug_key=aug_data` to augment data
Args:
data (tf.Tensor):
probe (bool, optional): Defaults to True.
aug_key (str, optional): Defaults to 'data'.
Returns:
dict: data_dict
"""
# creating local variable which will store copy of CTA log probabilities
with tf.variable_creator_scope(_skip_mirrored_creator):
local_log_prob = tf.Variable(
lambda: tf.ones(self.state_shape, dtype=tf.float32),
trainable=False,
name='cta_log_probs')
self.log_probs.append(local_log_prob)
output_dict = {}
if probe:
# 采样 [num_layers] 个 op_indices 和 op_args
probe_op_indices, probe_op_args = self._sample_ops_uniformly()
probe_data = self._apply_ops(data, probe_op_indices, probe_op_args)
output_dict['probe_op_indices'] = probe_op_indices
output_dict['probe_op_args'] = probe_op_args
output_dict['probe_data'] = probe_data
if aug_key is not None:
op_indices, op_args = self._sample_ops(local_log_prob)
aug_data = self._apply_ops(data,
op_indices,
op_args,
prob_to_apply=self.prob_to_apply)
output_dict[aug_key] = aug_data
if aug_key != 'data':
output_dict['data'] = data
return output_dict
def test_keras_layer_add_weight(self):
class Layer(base_layer.Layer):
def __init__(self):
super().__init__()
self.w = self.add_weight(
shape=(2, ),
initializer=lambda shape, dtype: [0, 1],
trainable=True)
self.b = self.add_weight(
shape=(2, ),
initializer=lambda shape, dtype: [2, 3],
trainable=False)
def sharded_variable_creator(next_creator, **kwargs):
v1_value = kwargs['initial_value']()[0:1]
v2_value = kwargs['initial_value']()[1:]
kwargs['initial_value'] = v1_value
kwargs['shape'] = (1, )
v1 = next_creator(**kwargs)
kwargs['initial_value'] = v2_value
kwargs['shape'] = (1, )
v2 = next_creator(**kwargs)
return sharded_variable.ShardedVariable([v1, v2])
with tf.variable_creator_scope(sharded_variable_creator):
layer = Layer()
self.assertLen(layer.trainable_weights, 2)
self.assertEqual(layer.trainable_weights[0], [0])
self.assertEqual(layer.trainable_weights[1], [1])
self.assertLen(layer.non_trainable_weights, 2)
self.assertEqual(layer.non_trainable_weights[0], [2])
self.assertEqual(layer.non_trainable_weights[1], [3])
self.assertAllEqual(
layer.weights,
layer.trainable_weights + layer.non_trainable_weights)
self.assertAllEqual(layer.trainable_weights, layer.trainable_variables)
self.assertAllEqual(layer.weights, layer.variables)
checkpoint_deps = set(dep.ref
for dep in layer._checkpoint_dependencies)
self.assertEqual(checkpoint_deps, set([layer.w, layer.b]))
def network_initializer(self):
"""
"""
with tf.variable_creator_scope('convolution_network') as scope:
output = self.
def initialize():
with tf.variable_creator_scope(variable_utils.create_tensor_variable):
return tf.nest.map_structure(lambda m: m.variables,
metrics_constructor())
def finalize(state):
with tf.variable_creator_scope(
build_replace_variable_with_parameter_creator(state)):
return tf.nest.map_structure(lambda metric: metric.result(),
metrics_constructor())
def trace_and_update_module(
module: tf.Module, tf_function: function.Function, name: str,
strip_control_dependencies: bool) -> function.ConcreteFunction:
"""Traces `tf_function` and saves under attr `name` of `module`.
Args:
module: A saveable module which will contain the traced `tf_function` under
attr `name`.
tf_function: A tf.function to trace.
name: A name to same the traced `tf_function` to.
strip_control_dependencies: Boolean. If True, automatic control dependencies
will be stripped from the outputs of `tf_function`. This should almost
always be False. It is useful only if you want to use the structure of the
TF graph to perform any graph manipulations.
Returns:
The concrete function obtained from tracing `tf_function`.
"""
resource_tracker = tracking.ResourceTracker()
object_tracker = annotators.ObjectTracker()
created_variables = []
def _variable_creator(next_creator, **kwargs):
var = next_creator(**kwargs)
created_variables.append(var)
return var
# Trace `tf_function` to gather any resources in it using the
# resource_tracker. These are then assigned to `module.resources` and tracked
# before exporting to SavedModel.
with tracking.resource_tracker_scope(resource_tracker), \
annotators.object_tracker_scope(object_tracker), \
tf.variable_creator_scope(_variable_creator):
concrete_fn = tf_function.get_concrete_function()
# Prior to 2020/10/08, saving a tf.function with a concrete function signature
# would ensure that the function was not re-traced in a round-trip to a
# SavedModel. Since this is no longer the case, we save the concrete function
# directly.
if tf.compat.forward_compatible(2020, 10, 8):
pruned_function = optimize_concrete_function(concrete_fn,
strip_control_dependencies)
module.pruned_variables = pruned_function.variables
setattr(module, name, pruned_function)
else:
setattr(module, name, tf_function)
# Any variables created need to be explicitly tracked.
module.created_variables = created_variables
# Resources need to be explicitly tracked.
module.resources = resource_tracker.resources
module.trackable_objects = object_tracker.trackable_objects
# TODO(b/158011374) - Stop explicitly tracking initializers. Tracking the
# table should be sufficient.
initializers = []
for resource in module.resources:
if isinstance(resource, lookup_ops.InitializableLookupTableBase):
initializers.append(resource._initializer) # pylint: disable=protected-access
module.initializers = initializers
module.assets = [
common_types.Asset(asset_filepath) for asset_filepath in
concrete_fn.graph.get_collection(tf.compat.v1.GraphKeys.ASSET_FILEPATHS)
]
return concrete_fn
def trace_and_write_v2_saved_model(saved_model_dir, preprocessing_fn,
input_signature, base_temp_dir,
tensor_replacement_map,
output_keys_to_name_map):
"""Writes out a SavedModelV2 with preprocessing_fn traced using tf.function.
The SavedModel written contains a method called `transform_fn` that
represents the traced `preprocessing_fn`. Additionally, if this is the final
SavedModel being written out, it will contain a method called `metadata_fn`
that provides deferred schema annotations.
Args:
saved_model_dir: Path to write SavedModel to.
preprocessing_fn: A user defined python function to be traced.
input_signature: TypeSpecs describing the inputs to the `preprocessing_fn`.
base_temp_dir: Base path to write temporary artifacts to.
tensor_replacement_map: A map from placeholder tensor names to their
evaluated replacement tensors.
output_keys_to_name_map: A map from output dictionary keys to the names of
the tensors that they represent.
Returns:
A tuple containing a pair of `tf.ConcreteFunction`s:
1. The traced preprocessing_fn.
2. A metadata_fn that returns a dictionary containing the deferred
annotations added to the graph when invoked with any valid input.
"""
module = tf.Module()
transform_fn = get_traced_transform_fn(
preprocessing_fn,
input_signature,
base_temp_dir,
tensor_replacement_map=tensor_replacement_map,
output_keys_to_name_map=output_keys_to_name_map)
metadata_fn = None
resource_tracker = tracking.ResourceTracker()
created_variables = []
def _variable_creator(next_creator, **kwargs):
var = next_creator(**kwargs)
created_variables.append(var)
return var
# TODO(b/164921571): Handle generic Trackable objects.
# Trace the `transform_fn` and `metadata_fn` to gather any resources in it
# using the resource_tracker. These are then assigned to `module.resources`
# and tracked before exporting to SavedModel.
with tracking.resource_tracker_scope(
resource_tracker), tf.variable_creator_scope(_variable_creator):
concrete_transform_fn = transform_fn.get_concrete_function()
concrete_metadata_fn = None
# If the `TENSOR_REPLACEMENTS` graph collection is empty, all TFT analyzers
# in the `preprocessing_fn` have already been evaluated.
if not concrete_transform_fn.graph.get_collection(
analyzer_nodes.TENSOR_REPLACEMENTS):
metadata_fn = schema_inference.get_traced_metadata_fn(
tensor_replacement_map,
preprocessing_fn,
input_signature,
base_temp_dir,
evaluate_schema_overrides=True)
concrete_metadata_fn = metadata_fn.get_concrete_function()
# Save ConcreteFunction when possible since the above workaround won't work if
# the tf.function is retraced.
if tf.compat.forward_compatible(2020, 10, 8):
module.transform_fn = concrete_transform_fn
module.metadata_fn = concrete_metadata_fn
else:
module.transform_fn = transform_fn
module.metadata_fn = metadata_fn
# Any variables created need to be explicitly tracked.
module.created_variables = created_variables
# Resources need to be explicitly tracked.
module.resources = resource_tracker.resources
# TODO(b/158011374) - Stop explicitly tracking initializers. Tracking the
# table should be sufficient.
initializers = []
for resource in module.resources:
if isinstance(resource, lookup_ops.InitializableLookupTableBase):
initializers.append(resource._initializer) # pylint: disable=protected-access
module.initializers = initializers
module.assets = [
common_types.Asset(asset_filepath)
for asset_filepath in concrete_transform_fn.graph.get_collection(
tf.compat.v1.GraphKeys.ASSET_FILEPATHS)
]
tf.saved_model.save(module, saved_model_dir)
return concrete_transform_fn, concrete_metadata_fn
def define_vggish(waveform):
with tf.variable_creator_scope(var_tracker):
features = waveform_to_features(waveform)
return vggish_slim.define_vggish_slim(features, training=False)
def __init__(self,
batch_size=64,
input_space=28 * 28,
latent_space=10,
p=3.,
middle_layers=None,
activation_fn=tf.nn.tanh,
learning_rate=0.001,
l2_lambda=0.001,
initializer_fn=he_initializer):
self.batch_size = batch_size
self.input_space = input_space
self.latent_space = latent_space
self.p = p
self.middle_layers = [1024, 1024]
self.activation_fn = activation_fn
self.learning_rate = learning_rate
self.initializer_fn = initializer_fn
tf.reset_default_graph()
self.input_x = tf.placeholder(tf.float32, [None, input_space])
self.z_tensor = tf.placeholder(tf.float32, [None, latent_space])
with tf.variable_scope('encoder'):
self._encoder()
self.encoded = self.encoder_layers[-1]
with tf.variable_scope('decoder'):
self.decoder_layers = self._decoder(self.encoded)
self.decoded = self.decoder_layers[-1]
tf.get_variable_scope().reuse_variables()
self.generator_layers = self._decoder(self.z_tensor)
self.generated = tf.nn.sigmoid(self.generator_layers[-1],
name='generated')
sizes = [64, 64, 1]
with tf.variable_scope('discriminator'):
self.disc_layers_neg = self._discriminator(self.encoded, sizes)
self.disc_neg = self.disc_layers_neg[-1]
tf.get_variable_scope().reuse_variables()
self.disc_layers_pos = self._discriminator(self.z_tensor, sizes)
self.disc_pos = self.disc_layers_pos[-1]
self.pos_loss = tf.nn.relu(self.disc_pos) - self.disc_pos + tf.log(
1.0 + tf.exp(-tf.abs(self.disc_pos)))
self.neg_loss = tf.nn.relu(self.disc_neg) - self.disc_neg + tf.log(
1.0 + tf.exp(-tf.abs(self.disc_neg)))
self.disc_loss = tf.reduce_mean(tf.add(self.pos_loss, self.neg_loss))
self.enc_loss = tf.reduce_mean(
tf.subtract(self.neg_loss, self.disc_neg))
batch_logloss = tf.reduce_sum(
tf.nn.sigmoid_cross_entropy_with_logits(logits=self.decoded,
labels=self.input_x), 1)
self.ae_loss = tf.reduce_mean(batch_logloss)
disc_ws = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES,
scope='discriminator')
ae_ws = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES,
scope='encoder') + tf.get_collection(
tf.GraphKeys.GLOBAL_VARIABLES,
scope='decoder')
self.l2_loss = tf.multiply(
tf.reduce_sum([tf.nn.l2_loss(ws) for ws in ae_ws]), l2_lambda)
self.gen_loss = tf.add(tf.add(self.enc_loss, self.ae_loss),
self.l2_loss)
with tf.variable_creator_scope('optimizer'):
self.train_discriminator = tf.train.RMSPropOptimizer(
self.learning_rate).minimize(self.disc_loss, var_list=disc_ws)
self.train_generator = tf.train.RMSPropOptimizer(
self.learning_rate).minimize(self.gen_loss, var_list=ae_ws)
self.sess = tf.Session()
init = tf.global_variables_initializer()
self.sess.run(init)
