def client_eval(incoming_model_weights, dataset):
"""Returns local outputs after evaluting `model_weights` on `dataset`."""
with tf.init_scope():
model = model_fn()
model_weights = model_utils.ModelWeights.from_model(model)
tf.nest.map_structure(lambda v, t: v.assign(t), model_weights,
incoming_model_weights)
def reduce_fn(num_examples, batch):
model_output = model.forward_pass(batch, training=False)
if model_output.num_examples is None:
# Compute shape from the size of the predictions if model didn't use the
# batch size.
return num_examples + tf.shape(model_output.predictions,
out_type=tf.int64)[0]
else:
return num_examples + tf.cast(model_output.num_examples,
tf.int64)
dataset_reduce_fn = dataset_reduce.build_dataset_reduce_fn(
use_experimental_simulation_loop)
num_examples = dataset_reduce_fn(
reduce_fn=reduce_fn,
dataset=dataset,
initial_state_fn=lambda: tf.zeros([], dtype=tf.int64))
return collections.OrderedDict(
local_outputs=model.report_local_outputs(),
num_examples=num_examples)
def build(self, input_shape):
if tf.__version__.startswith('1.5'):
self.mean = self.add_weight(name='mean',
shape=(self.b, input_shape[1]),
initializer=tf.constant_initializer(
np.tile(self.m, (self.b, 1)),
verify_shape=True),
trainable=False)
self.std = self.add_weight(name='std',
shape=(self.b, input_shape[1]),
initializer=tf.constant_initializer(
np.tile(self.s, (self.b, 1))),
trainable=False)
elif tf.__version__.startswith('2.'):
with tf.init_scope():
self.mean = self.add_weight(
name='mean',
shape=(self.b, input_shape[1]),
#initializer=tf.constant_initializer(np.tile(self.m,(self.b,1)),verify_shape=True),
initializer=tf.constant_initializer(
np.tile(self.m, (self.b, 1))),
trainable=False)
self.std = self.add_weight(
name='std',
shape=(self.b, input_shape[1]),
#initializer=tf.constant_initializer(np.tile(self.s,(self.b,1)),verify_shape=True),
initializer=tf.constant_initializer(
np.tile(self.s, (self.b, 1))),
trainable=False)
else:
sys.exit("Tensforflow version " + tf.__version__ +
" unknown for preprocessing layer")
super(PreprocessLayer,
self).build(input_shape) # Be sure to call this at the end
def __call__(self, model, data):
x, y = data
batch_size = tf.shape(x)[0]
with tf.GradientTape() as tape:
y_pred = model(x, training=True)
loss = model.compiled_loss(y, y_pred, regularization_losses=model.losses)
gradients = tape.gradient(loss, model.trainable_weights)
gradients = [g * tf.cast(batch_size, tf.float32) for g in gradients]
sample = self.normal_initializer(tf.shape(y_pred))# * tf.sqrt(tf.cast(y_pred, tf.float32))
sample = model.sample_observations(y_pred, sample)
sample = tf.reduce_sum(sample, 0, keepdims=True)
single_x = tf.expand_dims(x[0,], axis=0)
with tf.GradientTape() as tape:
observations = model(single_x, training=True)
fim_samples = tape.gradient(observations, model.trainable_weights, output_gradients=sample)
fim_samples = [tf.expand_dims(samples, 0) for samples in fim_samples]
fim_diagonal = model.fim_diagonal(x)
with tf.init_scope():
model.optimizer._create_slots(model.trainable_weight_initializers())
model.optimizer.apply_information(zip(fim_samples, fim_diagonal, model.trainable_weights))
model.optimizer.apply_gradients(zip(gradients, model.trainable_weights))
model.compiled_metrics.update_state(y, y_pred)
# Return a dict mapping metric names to current value
return {m.name: m.result() for m in model.metrics}
def server_update(global_model, mean_model_delta, optimizer_state):
"""Updates the global model with the mean model update from clients."""
with tf.init_scope():
# Create a structure of variables that the server optimizer can update.
model_variables = tf.nest.map_structure(
lambda t: tf.Variable(initial_value=tf.zeros(t.shape, t.dtype)
), global_model)
optimizer = keras_optimizer.build_or_verify_tff_optimizer(
server_optimizer_fn,
model_variables.trainable,
disjoint_init_and_next=True)
# Set the variables to the current global model, the optimizer will
# update these variables.
tf.nest.map_structure(lambda a, b: a.assign(b), model_variables,
global_model)
# We might have a NaN value e.g. if all of the clients processed had no
# data, so the denominator in the federated_mean is zero. If we see any
# NaNs, zero out the whole update.
# TODO(b/124538167): We should increment a server counter to
# track the fact a non-finite weights_delta was encountered.
finite_weights_delta, _ = tensor_utils.zero_all_if_any_non_finite(
mean_model_delta)
# Update the global model variables with the delta as a pseudo-gradient.
negative_weights_delta = tf.nest.map_structure(lambda w: -1.0 * w,
finite_weights_delta)
optimizer_state, updated_weights = optimizer.next(
optimizer_state, model_variables.trainable, negative_weights_delta)
# Keras optimizers mutate model variables in with the `next` step above, so
# we skip calling the assignment for those optimizers.
if not isinstance(optimizer, keras_optimizer.KerasOptimizer):
tf.nest.map_structure(lambda a, b: a.assign(b),
model_variables.trainable, updated_weights)
return model_variables, optimizer_state
def call(self, inputs, training=None, **kwargs):
with tf.init_scope():
if not self._is_set:
self.layer(inputs)
self._is_set = True
if training:
w = self.layer.kernel
w_shape = w.shape.as_list()
w = tf.reshape(w, (-1, self.d))
u_hat = self.u
v_hat = None
for _ in range(self.power_iteration):
v_ = tf.matmul(u_hat, w, transpose_b=True)
v_hat = self.l2_normalize(v_)
u_ = tf.matmul(v_hat, w)
u_hat = self.l2_normalize(u_)
# u_hat = tf.stop_gradient(u_hat)
# v_hat = tf.stop_gradient(v_hat)
sigma = tf.matmul(tf.matmul(v_hat, w), u_hat, transpose_b=True)
sigma = tf.reshape(sigma, ())
self.u.assign(u_hat)
self.sigma.assign(sigma)
return self.layer(inputs, **kwargs) / (self.sigma + K.epsilon())
def update(dataset: tf.data.Dataset, state: State, message: server.Message,
model_fn: Callable, optimizer_fn: Callable) -> Output:
with tf.init_scope():
model = model_fn(pos_weight=state.client_pos_weight)
optimizer = optimizer_fn()
message.model.assign_weights_to(model)
def training_fn(num_examples, batch):
with tf.GradientTape() as tape:
outputs = model.forward_pass(batch, training=True)
optimizer.apply_gradients(
zip(tape.gradient(outputs.loss, model.trainable_variables),
model.trainable_variables))
return num_examples + outputs.num_examples
client_weight = dataset.reduce(tf.constant(0, dtype=tf.int32), training_fn)
weights_delta = tf.nest.map_structure(lambda a, b: a - b,
model.trainable_variables,
message.model.trainable)
return Output(weights_delta=weights_delta,
metrics=model.report_local_outputs(),
client_weight=tf.cast(client_weight, dtype=tf.float32),
client_state=State(
client_index=state.client_index,
client_pos_weight=state.client_pos_weight))
def set_optimizer_params(optimizer, param_list):
# Set optimizer tracking parameters
if len(param_list) > 0:
with tf.name_scope(optimizer._name):
with tf.init_scope():
optimizer._create_all_weights(param_list)
return optimizer
def __call__(self, model, data):
x, y = data
batch_size = tf.shape(x)[0]
with tf.GradientTape() as tape:
y_pred = model(x, training=True)
loss = model.compiled_loss(y, y_pred, regularization_losses=model.losses)
gradients = tape.gradient(loss, model.trainable_weights)
gradients = [g * tf.cast(batch_size, tf.float32) for g in gradients]
sample = model.sample_observations(y_pred, self.normal_initializer(tf.shape(y_pred)))
def sample_fim(x_and_sample):
with tf.GradientTape() as tape:
y_pred = model(tf.expand_dims(x_and_sample[0], 0), training=True)
return tape.gradient(y_pred, model.trainable_weights, output_gradients=tf.expand_dims(x_and_sample[1], 0))
fim_samples = tf.vectorized_map(sample_fim, (x, sample))
fim_diagonal = model.fim_diagonal(x)
with tf.init_scope():
model.optimizer._create_slots(model.trainable_weight_initializers())
model.optimizer.apply_information(zip(fim_samples, fim_diagonal, model.trainable_weights))
model.optimizer.apply_gradients(zip(gradients, model.trainable_weights))
model.compiled_metrics.update_state(y, y_pred)
# Return a dict mapping metric names to current value
return {m.name: m.result() for m in model.metrics}
def _create_special_tokens_dict(self, vocab_table, vocab_file):
special_tokens = dict(start_of_sequence_id="[CLS]",
end_of_segment_id="[SEP]",
padding_id="[PAD]",
mask_id="[MASK]")
with tf.init_scope():
if tf.executing_eagerly():
special_token_ids = vocab_table.lookup(
tf.constant(list(special_tokens.values()), tf.string))
else:
# A blast from the past: non-eager init context while building Model.
# This can happen with Estimator or tf.compat.v1.disable_v2_behavior().
logging.warning(
"Non-eager init context; computing "
"BertTokenizer's special_tokens_dict in tf.compat.v1.Session"
)
with tf.Graph().as_default():
local_vocab_table, _ = self._create_vocab_table_and_initializer(
vocab_file)
special_token_ids_tensor = local_vocab_table.lookup(
tf.constant(list(special_tokens.values()), tf.string))
init_ops = [tf.compat.v1.initialize_all_tables()]
with tf.compat.v1.Session() as sess:
sess.run(init_ops)
special_token_ids = sess.run(special_token_ids_tensor)
result = dict()
for k, v in zip(special_tokens, special_token_ids):
v = int(v) # Numpy to Python.
if v >= 0:
result[k] = v
else:
logging.warning(
"Could not find %s as token \"%s\" in vocab file %s",
k, special_tokens[k], vocab_file)
return result
def __init__(
self, values_shape=(), values_dtype="int8", name="flip_ratio", dtype=None
):
super().__init__(name=name, dtype=dtype)
self.values_dtype = tf.as_dtype(values_dtype)
self.values_shape = tf.TensorShape(values_shape).as_list()
self.is_weight_metric = True
with tf.init_scope():
self._previous_values = self.add_weight(
"previous_values",
shape=values_shape,
dtype=self.values_dtype,
initializer=tf.keras.initializers.zeros,
aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA,
)
self.total = self.add_weight(
"total",
initializer=tf.keras.initializers.zeros,
aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA,
)
self.count = self.add_weight(
"count",
initializer=tf.keras.initializers.zeros,
aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA,
)
self._size = np.prod(self.values_shape)
def get_summary_writer(save_dir, subdir='', comm=MPI.COMM_WORLD):
if comm.Get_rank() != 0:
return None
if save_dir is None:
return None
with tf.init_scope():
return summary.create_file_writer(os.path.join(save_dir, 'tb', subdir))
def local_update(initial_weights, data):
# TODO(b/190334722): Restructure so that model_fn only needs to be invoked
# once.
with tf.init_scope():
model = model_fn()
model_weights = model_utils.ModelWeights.from_model(model)
tf.nest.map_structure(lambda weight, value: weight.assign(value),
model_weights, initial_weights)
num_examples = tf.constant(0, tf.int32)
optimizer_state = optimizer.initialize(optimizer_tensor_specs)
# TODO(b/161529310): Different from creating an iterator using iter(data).
for batch in data:
with tf.GradientTape() as tape:
outputs = model.forward_pass(batch)
gradients = tape.gradient(outputs.loss, model_weights.trainable)
num_examples += tf.shape(outputs.predictions)[0]
optimizer_state, updated_weights = optimizer.next(
optimizer_state, _flat_tuple(model_weights.trainable),
_flat_tuple(gradients))
updated_weights = tf.nest.pack_sequence_as(model_weights.trainable,
updated_weights)
tf.nest.map_structure(lambda weight, value: weight.assign(value),
model_weights.trainable, updated_weights)
model_delta = tf.nest.map_structure(lambda x, y: x - y,
initial_weights.trainable,
model_weights.trainable)
return ClientResult(
update=model_delta, update_weight=tf.cast(num_examples, tf.float32))
def _build_from_signature(self, query, value, key=None):
super(MultiHeadRelativeAttention,
self)._build_from_signature(query=query, value=value, key=key)
if hasattr(value, "shape"):
value_shape = tf.TensorShape(value.shape)
else:
value_shape = value
if key is None:
key_shape = value_shape
elif hasattr(key, "shape"):
key_shape = tf.TensorShape(key.shape)
else:
key_shape = key
common_kwargs = dict(kernel_initializer=self._kernel_initializer,
bias_initializer=self._bias_initializer,
kernel_regularizer=self._kernel_regularizer,
bias_regularizer=self._bias_regularizer,
activity_regularizer=self._activity_regularizer,
kernel_constraint=self._kernel_constraint,
bias_constraint=self._bias_constraint)
with tf.init_scope():
einsum_equation, _, output_rank = _build_proj_equation(
key_shape.rank - 1, bound_dims=1, output_dims=2)
self._encoding_dense = tf.keras.layers.EinsumDense(
einsum_equation,
output_shape=_get_output_shape(
output_rank - 1, [self._num_heads, self._key_dim]),
bias_axes=None,
name="encoding",
**common_kwargs)
def _set_initializers(self):
"""Change initializers to load a model from a tensorflow checkpoint."""
if self.comm.Get_rank() > 0 or self.train_dir == 'test':
return
assert self.model.built
checkpoint_scope = 'reward_model'
with tf.init_scope():
# Initialize!
params = {v.op.name: v for v in self.get_params()}
checkpoint = tf.train.latest_checkpoint(
os.path.join(self.train_dir, 'checkpoints/'))
available = tf.train.list_variables(checkpoint)
unchanged = {}
for name, shape in available:
if not name.startswith(checkpoint_scope + '/'):
# print('skipping', name)
continue
if name.endswith('adam') or name.endswith('adam_1'):
# print('skipping', name)
continue
print('setting', name)
var = params[self.scope + name[len(checkpoint_scope):]]
assert var.shape == shape, 'Shape mismatch: %s.shape = %s != %s' % (
var.op.name, var.shape, shape)
unchanged[name] = var
tf.train.init_from_checkpoint(checkpoint, unchanged)
def server_update(global_model, mean_model_delta, optimizer_state):
"""Updates the global model with the mean model update from clients."""
with tf.init_scope():
model = model_fn()
optimizer = server_optimizer_fn()
# We must force variable creation for momentum and adaptive optimizers.
_eagerly_create_optimizer_variables(model=model,
optimizer=optimizer)
model_variables = model_utils.ModelWeights.from_model(model)
optimizer_variables = optimizer.variables()
# Set the variables to the current global model, the optimizer will
# update these variables.
tf.nest.map_structure(lambda a, b: a.assign(b),
(model_variables, optimizer_variables),
(global_model, optimizer_state))
# We might have a NaN value e.g. if all of the clients processed had no
# data, so the denominator in the federated_mean is zero. If we see any
# NaNs, zero out the whole update.
# TODO(b/124538167): We should increment a server counter to
# track the fact a non-finite weights_delta was encountered.
finite_weights_delta, _ = tensor_utils.zero_all_if_any_non_finite(
mean_model_delta)
# Update the global model variables with the delta as a pseudo-gradient.
_apply_delta(optimizer=optimizer,
model=model,
delta=finite_weights_delta)
return model_variables, optimizer_variables
def __init__(self, spec, meta_graph, trainable, checkpoint_path, name):
"""Private constructor.
Args:
spec: _ModuleSpec instance.
meta_graph: MetaGraphDef to use
trainable: whether module is trainable.
checkpoint_path: None or a string to the variables checkpoints.
name: variable and scope name where to instantiate the Module. Must be an
unused name scope.
"""
self._spec = spec
self._meta_graph = meta_graph
self._trainable = trainable
self._checkpoint_path = checkpoint_path
register_ops_if_needed({
op.name
for op in self._meta_graph.meta_info_def.stripped_op_list.op
})
# Use an init scope to clear dependencies and lift state ops outside of
# function-building graphs. Modules can be constructed from deep inside
# functions that have dependencies active. Note that the dependencies
# would be active when applying the Module signature, just not active
# when creating the Module state. This use case has showed up in some
# TPU training code.
with tf.init_scope():
self._init_state(name)
def _compute_rolled_dates_table(self, roll_convention):
"""Computes and caches rolled dates table."""
already_computed = self._table_cache.rolled_dates.get(
roll_convention, None)
if already_computed is not None:
return already_computed
roll_convention_np = _to_np_roll_convention(roll_convention)
holidays_arg = self._holidays_np
if holidays_arg is None:
holidays_arg = [] # np.busday_offset doesn't accept None
adjusted_np = np.busday_offset(dates=self._dates_np,
offsets=0,
roll=roll_convention_np,
weekmask=1 - self._weekend_mask,
holidays=holidays_arg)
rolled_date_table = adjusted_np.astype(np.int32) + _ORDINAL_OF_1_1_1970
# To make tensor caching safe, lift the ops out of the current scope using
# tf.init_scope(). This allows e.g. to cache these tensors in one
# tf.function and reuse them in another tf.function.
with tf.init_scope():
rolled_date_table = tf.convert_to_tensor(rolled_date_table,
name="rolled_date_table")
self._table_cache.rolled_dates[roll_convention] = rolled_date_table
return rolled_date_table
def _get_temporary_graph_tensor(tensor_info, name):
"""Get a temporary graph tensor using attributes in `tensor_info`."""
is_asset_filepath = tensor_info.temporary_asset_value is not None
if is_asset_filepath:
# Placeholders cannot be used for assets as they will be initialized as part
# of the init op. Hence, a temporary file is written out during tracing.
# TODO(b/164921571) Support temporary files in tfrecord format.
# TODO(b/149997088): Reduce number of temporary files written out.
with tf.init_scope():
temporary_asset_filepath = os.path.join(
TFGraphContext.get_or_create_temp_dir(),
uuid.uuid4().hex)
with tf.io.gfile.GFile(temporary_asset_filepath, 'w') as f:
f.write(tensor_info.temporary_asset_value)
result = tf.constant(
temporary_asset_filepath,
dtype=tensor_info.dtype,
shape=tensor_info.shape,
name=name)
else:
# Using a placeholder with no default value causes tracing to fail if there
# is any control flow dependent on a child tensor of this placeholder.
# Hence, provide a temporary default value for it.
# If dtype is string, we want a tensor that contains '0's instead of b'[] to
# allow string to numeric conversion ops to trace successfully.
temporary_dtype = (
tf.int64 if tensor_info.dtype == tf.string else tensor_info.dtype)
temporary_tensor = tf2_utils.supply_missing_tensor(
1, tf.TensorShape(tensor_info.shape), temporary_dtype)
if tensor_info.dtype == tf.string:
temporary_tensor = tf.strings.as_string(temporary_tensor)
result = tf.raw_ops.PlaceholderWithDefault(
input=temporary_tensor, shape=tensor_info.shape)
return result
def is_xla_available():
"""Is XLA compilation available for the current device context?"""
global _IS_XLA_AVAILABLE
# There's unfortunately no cleaner way to get the device other than creating a
# new op and querying it.
with tf.name_scope("is_xla_available"):
device = tf.constant(0.0).device
if device not in _IS_XLA_AVAILABLE:
try:
# Take ourselves outside of any tf.function calls.
with tf.init_scope():
# Create temporary xla subgraph
with tf.compat.v1.Graph().as_default():
# We'll use a session so we can be compatible with both TF1 and TF2
with tf.compat.v1.Session() as sess:
# Check for XLA on the given device.
with tf.device(device):
sess.run(
tf.xla.experimental.compile(
lambda: tf.constant(0.0)))
except (ValueError, tf.errors.InvalidArgumentError):
_IS_XLA_AVAILABLE[device] = False
else:
_IS_XLA_AVAILABLE[device] = True
return _IS_XLA_AVAILABLE[device]
def _get_temporary_analyzer_output(
temp_dir: str,
tensor_info: TensorInfo,
name: Optional[str] = None) -> TemporaryAnalyzerOutputWrapper:
"""Create a temporary graph tensor using attributes in `tensor_info`.
Args:
temp_dir: Path to a directory to write out any temporary asset files to.
tensor_info: A `TensorInfo` object containing attributes to create the graph
tensor.
name: A string (or None). The created graph tensor uses this name.
Returns:
A named tuple `TemporaryAnalyzerOutputWrapper` with:
asset: If the graph tensor represents a path to an asset file, a
`tf.saved_model.Asset` object for tracking. Else, None.
graph_tensor: The graph tensor
"""
asset = None
with tf.name_scope('temporary_analyzer_output'):
is_asset_filepath = tensor_info.temporary_asset_value is not None
if is_asset_filepath:
# Placeholders cannot be used for assets, if this graph will be serialized
# to a SavedModel, as they will be initialized with the init op. If a
# `temp_dir` is provided, it is assumed that this graph will be
# serialized and a temporary asset file is written out . Else, a
# placeholder is returned.
# TODO(b/164921571) Support temporary files in tfrecord format.
# TODO(b/149997088): Reduce number of temporary files written out.
if temp_dir:
with tf.init_scope():
temporary_asset_filepath = os.path.join(temp_dir, uuid.uuid4().hex)
with tf.io.gfile.GFile(temporary_asset_filepath, 'w') as f:
f.write(tensor_info.temporary_asset_value)
# Wrap asset files using `tf.saved_model.Asset` to ensure that
# `SavedModel`s exported are hermetic.
asset = common_types.Asset(temporary_asset_filepath)
graph_tensor = tf.constant(
temporary_asset_filepath,
dtype=tensor_info.dtype,
shape=tensor_info.shape,
name=name)
else:
graph_tensor = tf.raw_ops.Placeholder(
dtype=tensor_info.dtype, shape=tensor_info.shape, name=name)
else:
# Using a placeholder with no default value causes tracing to fail if
# there is any control flow dependent on a child tensor of this
# placeholder. Hence, provide a temporary default value for it.
# If dtype is string, we want a tensor that contains '0's instead of b'[]
# to allow string to numeric conversion ops to trace successfully.
temporary_dtype = (
tf.int64 if tensor_info.dtype == tf.string else tensor_info.dtype)
temporary_tensor = tf2_utils.supply_missing_tensor(
1, tf.TensorShape(tensor_info.shape), temporary_dtype)
if tensor_info.dtype == tf.string:
temporary_tensor = tf.strings.as_string(temporary_tensor)
graph_tensor = tf.raw_ops.PlaceholderWithDefault(
input=temporary_tensor, shape=tensor_info.shape, name=name)
return TemporaryAnalyzerOutputWrapper(asset, graph_tensor)
def _check_same_graph(self):
"""Checks that the module is not being connect to multiple Graphs.
An instance of a Sonnet module 'owns' the variables it contains, and permits
seamless variable sharing. As such, connecting a single module instance to
multiple Graphs is not possible - this function will raise an error should
that occur.
Raises:
DifferentGraphError: if the module is connected to a different Graph than
it was previously used in.
"""
with tf.init_scope():
# We need `init_scope` incase we're running inside a defun. In that case
# what we want is information about where the function will be called not
# where the function is being built.
current_graph = tf.get_default_graph()
will_call_in_eager_context = tf.executing_eagerly()
if self._graph is None:
self._graph = current_graph
self._set_module_info()
if not will_call_in_eager_context:
# Same graph checks only make sense when calling from graph mode (in eager
# mode there is a single process level context where all modules are
# created).
if self._graph != current_graph:
raise DifferentGraphError(
"Cannot connect module to multiple Graphs.")
def apply_gradients(self,
grads_and_vars,
name: Optional[str] = None,
**kwargs):
"""Apply gradients to variables for each optimizer.
On the first call to `apply_gradients()`, compute the mapping from variables to
optimizers and cache it in the `self.var_opt_mapping` dict for serialization and
faster access.
"""
if self.var_opt_mapping is None:
# Convert `grads_and_vars` to list so we can iterate multiple times over it
grads_and_vars = list(grads_and_vars)
self._compute_var_opt_mapping(grads_and_vars)
# Split gradients and variables into a separate list for each optimizer
grad_var_lists = [[] for _ in range(len(self.pred_opt_pairs) + 1)]
for grad, var in grads_and_vars:
if var.name in self.var_opt_mapping:
grad_var_lists[self.var_opt_mapping[var.name]].append(
(grad, var))
with tf.init_scope():
for optimizer, opt_grads_and_vars in zip(self.optimizers,
grad_var_lists):
optimizer._create_slots([v for (_, v) in grads_and_vars])
return tf.distribute.get_replica_context().merge_call(
self._apply_gradients, args=(grad_var_lists, name), kwargs=kwargs)
def _get_step(self):
with tf.init_scope():
if context.executing_eagerly():
graph = None
else:
graph = tf.get_default_graph()
return self._get_non_slot_variable('step', graph=graph)
def finalizer(unfinalized_metric_values: List[tf.Tensor]):
# Construct a new keras metirc here, which is necessary because this
# `tf.function` may be invoked in a different context as the `model_fn`, and
# we need the `tf.Variable`s to be created in the current scope in order to
# use `keras_metric.result()`.
with tf.init_scope():
keras_metric = create_keras_metric(metric)
py_typecheck.check_type(unfinalized_metric_values, list)
if len(keras_metric.variables) != len(unfinalized_metric_values):
raise ValueError(
'The input to the finalizer should be a list of `tf.Tensor`s matching'
f' the variables of the Keras metric {keras_metric.name}. Expected '
f'a list of `tf.Tensor`s of length {len(keras_metric.variables)}, '
f'found a list of length {len(unfinalized_metric_values)}.')
for v, a in zip(keras_metric.variables, unfinalized_metric_values):
py_typecheck.check_type(a, tf.Tensor)
if v.shape != a.shape or v.dtype != a.dtype:
raise ValueError(
'The input to the finalizer should be a list of `tf.Tensor`s '
f'matching the variables of the Keras metric {keras_metric.name}. '
f'Expected a `tf.Tensor` of shape {v.shape} and dtype {v.dtype!r}, '
f'found a `tf.Tensor` of shape {a.shape} and dtype {a.dtype!r}.')
v.assign(a)
return keras_metric.result()
def __init__(
self,
active=True,
scale=1.,
print_loss=False,
print_batch_time=False,
return_lossval=False,
print_time=False, #compat, has no effect
**kwargs):
super(LossLayerBase, self).__init__(**kwargs)
if print_time:
print(
"print_time has no effect and is only for compatibility purposes"
)
self.active = active
self.scale = scale
self.print_loss = print_loss
self.print_batch_time = print_batch_time
self.return_lossval = return_lossval
with tf.init_scope():
now = tf.timestamp()
self.time = tf.Variable(-now,
name=self.name + '_time',
trainable=False)
def evaluate(dataset: tf.data.Dataset, state: State,
weights: tff.learning.ModelWeights, coefficient_fn: Callable,
model_fn: Callable) -> Evaluation:
with tf.init_scope():
mixing_coefficients = coefficient_fn()
model = model_fn(pos_weight=state.client_pos_weight)
tf.nest.map_structure(lambda v, t: v.assign(t), mixing_coefficients,
state.mixing_coefficients)
__mix_weights(mixing_coefficients, state.model,
weights).assign_weights_to(model)
def evaluation_fn(state, batch):
outputs = model.forward_pass(batch, training=False)
y_true = tf.reshape(batch[1], (-1, ))
y_pred = tf.round(
tf.nn.sigmoid(tf.reshape(outputs.predictions, (-1, ))))
return state + tf.math.confusion_matrix(y_true, y_pred, num_classes=2)
confusion_matrix = dataset.reduce(tf.zeros((2, 2), dtype=tf.int32),
evaluation_fn)
return Evaluation(confusion_matrix=confusion_matrix,
metrics=model.report_local_outputs())
def initialize(self, input_shape):
self.input_shape = input_shape
self.Dx_loss, self.Gx_loss = self.adv_loss_ctor(self.D_x)
self.Dy_loss, self.Gy_loss = self.adv_loss_ctor(self.D_y)
with tf.init_scope():
self.G_zx.initialize(input_shape)
self.G_zy.initialize(input_shape)
self._init_checkpoint()
def __init__(self, initializer):
super(Index, self).__init__()
if isinstance(initializer, trackable_base.Trackable):
self._initializer = self._track_trackable(initializer,
"_initializer")
with tf.init_scope():
self._resource_handle = self._create_resource()
self._init_op = self._initialize()
def _generate(self, feature_map_shape_list):
"""Generates a collection of bounding boxes to be used as anchors.
Args:
feature_map_shape_list: list of pairs of convnet layer resolutions in the
format [(height_0, width_0)]. For example, setting
feature_map_shape_list=[(8, 8)] asks for anchors that correspond
to an 8x8 layer. For this anchor generator, only lists of length 1 are
allowed.
Returns:
boxes_list: a list of BoxLists each holding anchor boxes corresponding to
the input feature map shapes.
Raises:
ValueError: if feature_map_shape_list, box_specs_list do not have the same
length.
ValueError: if feature_map_shape_list does not consist of pairs of
integers
"""
if not (isinstance(feature_map_shape_list, list)
and len(feature_map_shape_list) == 1):
raise ValueError(
'feature_map_shape_list must be a list of length 1.')
if not all([
isinstance(list_item, tuple) and len(list_item) == 2
for list_item in feature_map_shape_list
]):
raise ValueError('feature_map_shape_list must be a list of pairs.')
# Create constants in init_scope so they can be created in tf.functions
# and accessed from outside of the function.
with tf.init_scope():
self._base_anchor_size = tf.cast(tf.convert_to_tensor(
self._base_anchor_size),
dtype=tf.float32)
self._anchor_stride = tf.cast(tf.convert_to_tensor(
self._anchor_stride),
dtype=tf.float32)
self._anchor_offset = tf.cast(tf.convert_to_tensor(
self._anchor_offset),
dtype=tf.float32)
grid_height, grid_width = feature_map_shape_list[0]
scales_grid, aspect_ratios_grid = ops.meshgrid(self._scales,
self._aspect_ratios)
scales_grid = tf.reshape(scales_grid, [-1])
aspect_ratios_grid = tf.reshape(aspect_ratios_grid, [-1])
anchors = tile_anchors(grid_height, grid_width, scales_grid,
aspect_ratios_grid, self._base_anchor_size,
self._anchor_stride, self._anchor_offset)
num_anchors = anchors.num_boxes_static()
if num_anchors is None:
num_anchors = anchors.num_boxes()
anchor_indices = tf.zeros([num_anchors])
anchors.add_field('feature_map_index', anchor_indices)
return [anchors]
def _create_optimizer(self):
"""Initializes the hyperparameters and sets the self._optimizer property."""
if self._optimizer:
return
if not self._layer_collection:
self.register_layers(self._model, self._loss)
if self._config['adapt_damping']:
if 'train_batch' not in self._kfac_kwargs:
raise ValueError(
'Must provide a train_batch tuple to use adaptive '
'damping. Use register_train_batch or pass it in '
'during optimizer construction.')
if 'loss_fn' not in self._kfac_kwargs:
self._kfac_kwargs['loss_fn'] = utils.get_loss_fn(
self._model,
self._loss,
loss_weights=self._config['loss_weights'])
with tf.name_scope(self._name):
with tf.init_scope():
# "iterations" property will create iterations if necessary.
_ = self.iterations
self._create_hypers()
self._kfac_kwargs.update(self._hyper)
try:
# We use the TF 1 variable_scope instead of the TF 2 recommended
# name_scope because we need to recover the variables created in this
# scope, which is not possible with name_scope.
with tf.variable_scope(self._tf_var_scope):
self._optimizer = _KFAC_OPT_CLASS(
layer_collection=self._layer_collection,
**self._kfac_kwargs)
except ValueError as e:
msg = str(e)
if re.search('Variable .* already exists', msg):
raise ValueError(
'You may have instantiated a KFAC Optimizer with the same name as '
'an existing one. Try resetting the default graph, instantiating '
'the optimizer with a different name, or changing the optimizer\'s '
'name.\nHere is the original ValueError:\n ' + msg)
elif re.search(
'Found the following errors with variable registration'
'.*gamma.*registered with wrong number of uses.*', msg):
# We don't regex the name batch_normalization because the user could
# have renamed the layer. We don't regex beta because they could have
# used BatchNorm without the shift.
raise ValueError(
'There may have been an issue registering BatchNormalization. Try '
'using tf.keras.backend.set_learning_phase before model '
'construction. An alternative solution is to use the unfused '
'batchnorm implementation (pass the argument fused=False to '
'BatchNormalization).\nHere is the original ValueError:\n '
+ msg)
else:
raise e
