def check_grads(grads_and_vars):
"""
Check wether the gradients contain Inf or Nan.
Args:
grads_and_vars:
list of tuple (grad, var),
normally the output of opt.compute_gradients
Output:
has_nan: bool, True if there is Nan, otherwise it will be False
amax: tensor denoting the maximum value in gradients
"""
has_nan_ops = []
amax_ops = []
for grad in grads_and_vars:
if isinstance(grad, tuple):
grad = grad[0]
if grad is not None:
if isinstance(grad, ops.IndexedSlices):
x = grad.values
else:
x = grad
if x.dtype != dtypes.float32:
x = math_ops.cast(x, dtypes.float32)
has_nan_ops.append(math_ops.reduce_any(math_ops.is_nan(x)))
amax_ops.append(math_ops.reduce_max(math_ops.abs(x)))
has_nan = math_ops.reduce_any(has_nan_ops)
amax = math_ops.reduce_max(amax_ops)
return has_nan, amax
def call(self, inputs):
self._called = True
if self._max_tokens is None:
out_depth = K.get_value(self.num_elements)
else:
out_depth = self._max_tokens
if self._output_mode == BINARY:
bool_one_hot_data = array_ops.one_hot(
inputs, depth=out_depth, on_value=True, off_value=False)
reduced_bool_data = math_ops.reduce_any(bool_one_hot_data, axis=1)
binary_data = math_ops.cast(reduced_bool_data, dtypes.int64)
binary_data.set_shape(tensor_shape.TensorShape((None, out_depth)))
return binary_data
one_hot_data = array_ops.one_hot(inputs, depth=out_depth)
counts = math_ops.reduce_sum(one_hot_data, axis=1)
if self._output_mode == COUNT:
count_data = math_ops.cast(counts, dtypes.int64)
count_data.set_shape(tensor_shape.TensorShape((None, out_depth)))
return count_data
tf_idf_data = math_ops.multiply(counts, self.tf_idf_weights)
tf_idf_data.set_shape(tensor_shape.TensorShape((None, out_depth)))
if self._output_mode == TFIDF:
return tf_idf_data
# We can only get here if we didn't recognize the passed mode.
raise ValueError("Unknown output mode %s" % self._output_mode)
def _get_next_as_optional(iterator, name=None):
"""Returns an empty dataset indicator and the next input from the iterator."""
replicas = []
worker_has_values = []
for i, worker in enumerate(iterator._input_workers.worker_devices): # pylint: disable=protected-access
if name is not None:
d = tf_device.DeviceSpec.from_string(worker)
new_name = "%s_%s_%d" % (name, d.job, d.task)
else:
new_name = None
with ops.device(worker):
worker_has_value, next_element = (
iterator._iterators[i].get_next_as_list(new_name)) # pylint: disable=protected-access
worker_has_values.append(worker_has_value)
# Make `replicas` a flat list of values across all replicas.
replicas.append(next_element)
# `global_has_value` indicates whether there is data in this global batch.
# We do a all-reduce across all the workers in the multi-worker case.
# TODO(b/126259107): Do strategy.reduce for CollectiveAllReduceStrategy.
if len(worker_has_values) > 1:
with ops.device(
iterator._input_workers.compute_devices_for_worker(0)[0]): # pylint: disable=protected-access
# Place the tf.reduce_any op in device 0 to minimize communication
# cost.
# TODO(b/128545270): Investigate why placing it on worker 0 will cause
# the entire data to copy back from device to host.
# TODO(anjalisridhar): Use strategy.reduce because you can reduce values
# in a single graph or multiple graphs using low level communication ops.
global_has_value = math_ops.reduce_any(worker_has_values)
else:
global_has_value = worker_has_values[0]
return global_has_value, replicas
def insert(self, ids, scores):
"""Insert the ids and scores into the TopN."""
with ops.control_dependencies(self.last_ops):
scatter_op = state_ops.scatter_update(self.id_to_score, ids,
scores)
larger_scores = math_ops.greater(scores, self.sl_scores[0])
def shortlist_insert():
larger_ids = array_ops.boolean_mask(math_ops.to_int64(ids),
larger_scores)
larger_score_values = array_ops.boolean_mask(
scores, larger_scores)
shortlist_ids, new_ids, new_scores = tensor_forest_ops.top_n_insert(
self.sl_ids, self.sl_scores, larger_ids,
larger_score_values)
u1 = state_ops.scatter_update(self.sl_ids, shortlist_ids,
new_ids)
u2 = state_ops.scatter_update(self.sl_scores, shortlist_ids,
new_scores)
return control_flow_ops.group(u1, u2)
# We only need to insert into the shortlist if there are any
# scores larger than the threshold.
cond_op = control_flow_ops.cond(math_ops.reduce_any(larger_scores),
shortlist_insert,
control_flow_ops.no_op)
with ops.control_dependencies([cond_op]):
self.last_ops = [scatter_op, cond_op]
def reduce_any(input_tensor, axis=None, keepdims=False):
"""A version of tf.reduce_any that eagerly evaluates if possible."""
v = get_static_value(input_tensor)
if v is None:
return math_ops.reduce_any(input_tensor, axis=axis, keepdims=keepdims)
else:
return v.any(axis=axis, keepdims=keepdims)
def next_inputs(self, time, outputs, state, sample_ids, name=None):
with tf.name_scope(name, "TacoScheduledOutputTrainingHelperNextInputs",
[time, outputs, state, sample_ids]):
(finished, base_next_inputs,
state) = (super(TacoScheduledOutputTrainingHelper,
self).next_inputs(
time=time,
outputs=outputs,
state=state,
sample_ids=sample_ids,
name=name,
))
sample_ids = math_ops.cast(sample_ids, tf.bool)
def maybe_sample():
"""Perform scheduled sampling."""
# Feed last output frame as next input. outputs is [N, output_dim * r]
last_output_frame = outputs[:, -self._output_dim:]
return array_ops.where(sample_ids, last_output_frame,
base_next_inputs)
all_finished = math_ops.reduce_all(finished)
no_samples = math_ops.logical_not(math_ops.reduce_any(sample_ids))
next_inputs = control_flow_ops.cond(
math_ops.logical_or(all_finished, no_samples),
lambda: base_next_inputs, maybe_sample)
return finished, next_inputs, state
def correlation(y_true, y_pred):
'''Pearson correlation coefficient. (metric)
The linear corrlation between y_true and y_pred is between -1.0 and 1.0, indicating
positive correlation and negative correlation respectively. In particular, if the
correlation is 0.0, it means y_true and y_pred are irrelevant linearly.
This function is implemented by:
corr = [mean(y_true * y_pred) - mean(y_true) * mean(y_pred)]
/ [ std(y_true) * std(m_y_pred) ]
This function has been revised to prevent the division fail (0/0). When either y_true
or y_pred is 0, the correlation would be set as 0.0.
Input:
y_true: label, tensor in any shape.
y_pred: prediction, tensor in any shape.
Output:
scalar, the mean linear correlation between y_true and y_pred.
'''
m_y_true = math_ops.reduce_mean(y_true, axis=0)
m_y_pred = math_ops.reduce_mean(y_pred, axis=0)
s_y_true = gen_math_ops.sqrt(math_ops.reduce_mean(gen_math_ops.square(y_true), axis=0) - gen_math_ops.square(m_y_true))
s_y_pred = gen_math_ops.sqrt(math_ops.reduce_mean(gen_math_ops.square(y_pred), axis=0) - gen_math_ops.square(m_y_pred))
s_denom = s_y_true * s_y_pred
s_numer = math_ops.reduce_mean(y_true * y_pred, axis=0) - m_y_true * m_y_pred
s_index = gen_math_ops.greater(s_denom, 0)
f1 = lambda: constant_op.constant(0.0)
f2 = lambda: math_ops.reduce_mean(array_ops.boolean_mask(s_numer,s_index)/array_ops.boolean_mask(s_denom,s_index))
return control_flow_ops.case([(math_ops.reduce_any(s_index), f2)], default=f1)
def testAxesType(self):
for dtype in [dtypes.int64, dtypes.int32]:
with self.test_session(use_gpu=True) as sess:
v = math_ops.reduce_any([True, True],
constant_op.constant(0, dtype=dtype))
tf_v = sess.run(v)
self.assertAllEqual(tf_v, True)
def next_inputs(self, time, outputs, state, sample_ids, name=None):
with ops.name_scope(name, "ScheduledOutputTrainingHelperNextInputs",
[time, outputs, state, sample_ids]):
(finished, base_next_inputs,
state) = (super(ScheduledOutputTrainingHelper,
self).next_inputs(time=time,
outputs=outputs,
state=state,
sample_ids=sample_ids,
name=name))
def maybe_sample():
"""Perform scheduled sampling."""
def maybe_concatenate_auxiliary_inputs(outputs_, indices=None):
"""Concatenate outputs with auxiliary inputs, if they exist."""
if self._auxiliary_input_tas is None:
return outputs_
next_time = time + 1
auxiliary_inputs = nest.map_structure(
lambda ta: ta.read(next_time),
self._auxiliary_input_tas)
if indices is not None:
auxiliary_inputs = array_ops.gather_nd(
auxiliary_inputs, indices)
return nest.map_structure(
lambda x, y: array_ops.concat((x, y), -1), outputs_,
auxiliary_inputs)
if self._next_input_layer is None:
return array_ops.where(
sample_ids,
maybe_concatenate_auxiliary_inputs(outputs),
base_next_inputs)
where_sampling = math_ops.cast(array_ops.where(sample_ids),
dtypes.int32)
where_not_sampling = math_ops.cast(
array_ops.where(math_ops.logical_not(sample_ids)),
dtypes.int32)
outputs_sampling = array_ops.gather_nd(outputs, where_sampling)
inputs_not_sampling = array_ops.gather_nd(
base_next_inputs, where_not_sampling)
sampled_next_inputs = maybe_concatenate_auxiliary_inputs(
self._next_input_layer(outputs_sampling), where_sampling)
base_shape = array_ops.shape(base_next_inputs)
return (array_ops.scatter_nd(indices=where_sampling,
updates=sampled_next_inputs,
shape=base_shape) +
array_ops.scatter_nd(indices=where_not_sampling,
updates=inputs_not_sampling,
shape=base_shape))
all_finished = math_ops.reduce_all(finished)
no_samples = math_ops.logical_not(math_ops.reduce_any(sample_ids))
next_inputs = control_flow_ops.cond(
math_ops.logical_or(all_finished, no_samples),
lambda: base_next_inputs, maybe_sample)
return (finished, next_inputs, state)
def get_cluster_assignment(pairwise_distances, centroid_ids):
"""Assign data points to the neareset centroids.
Tensorflow has numerical instability and doesn't always choose
the data point with theoretically zero distance as it's nearest neighbor.
Thus, for each centroid in centroid_ids, explicitly assign
the centroid itself as the nearest centroid.
This is done through the mask tensor and the constraint_vect tensor.
Args:
pairwise_distances: 2-D Tensor of pairwise distances.
centroid_ids: 1-D Tensor of centroid indices.
Returns:
y_fixed: 1-D tensor of cluster assignment.
"""
predictions = math_ops.argmin(
array_ops.gather(pairwise_distances, centroid_ids), dimension=0)
batch_size = array_ops.shape(pairwise_distances)[0]
# Deal with numerical instability
mask = math_ops.reduce_any(array_ops.one_hot(
centroid_ids, batch_size, True, False, axis=-1, dtype=dtypes.bool),
axis=0)
constraint_one_hot = math_ops.multiply(
array_ops.one_hot(centroid_ids,
batch_size,
array_ops.constant(1, dtype=dtypes.int64),
array_ops.constant(0, dtype=dtypes.int64),
axis=0,
dtype=dtypes.int64),
math_ops.to_int64(math_ops.range(array_ops.shape(centroid_ids)[0])))
constraint_vect = math_ops.reduce_sum(
array_ops.transpose(constraint_one_hot), axis=0)
y_fixed = array_ops.where(mask, constraint_vect, predictions)
return y_fixed
def terminate_when_all_zero(current_argument, residual_powers,
accumulator):
del current_argument, accumulator # not used for condition
do_exit = math_ops.reduce_any(
math_ops.greater(residual_powers,
array_ops.ones_like(residual_powers)))
return do_exit
def call(self, inputs):
self._called = True
if self._max_tokens is None:
out_depth = K.get_value(self.num_elements)
else:
out_depth = self._max_tokens
# If the input is a sparse tensor, we densify it with the default value of
# -1. Because -1 is ignored by one_hot, this effectively drops the non-set
# positions from the output encoding.
if isinstance(inputs, sparse_tensor.SparseTensor):
inputs = sparse_ops.sparse_tensor_to_dense(inputs, default_value=-1)
if self._output_mode == BINARY:
bool_one_hot_data = array_ops.one_hot(
inputs, depth=out_depth, on_value=True, off_value=False)
reduced_bool_data = math_ops.reduce_any(bool_one_hot_data, axis=1)
binary_data = math_ops.cast(reduced_bool_data, dtypes.int64)
binary_data.set_shape(tensor_shape.TensorShape((None, out_depth)))
return binary_data
one_hot_data = array_ops.one_hot(inputs, depth=out_depth)
counts = math_ops.reduce_sum(one_hot_data, axis=1)
if self._output_mode == COUNT:
count_data = math_ops.cast(counts, dtypes.int64)
count_data.set_shape(tensor_shape.TensorShape((None, out_depth)))
return count_data
tf_idf_data = math_ops.multiply(counts, self.tf_idf_weights)
tf_idf_data.set_shape(tensor_shape.TensorShape((None, out_depth)))
if self._output_mode == TFIDF:
return tf_idf_data
# We can only get here if we didn't recognize the passed mode.
raise ValueError("Unknown output mode %s" % self._output_mode)
def testAxesType(self):
for dtype in [dtypes.int64, dtypes.int32]:
with self.session(use_gpu=True) as sess:
v = math_ops.reduce_any([True, True],
constant_op.constant(0, dtype=dtype))
tf_v = self.evaluate(v)
self.assertAllEqual(tf_v, True)
def call(self, inputs):
inputs = self._preprocess(inputs)
# If we're not doing any output processing, return right away.
if self._output_mode is None:
return inputs
# The table lookup ops don't natively support ragged tensors, so if we have
# a RT we need to use map_flat_values to look up every element.
if ragged_tensor.is_ragged(inputs):
indexed_data = ragged_functional_ops.map_flat_values(
self._table.lookup, inputs)
else:
indexed_data = self._table.lookup(inputs)
if self._output_mode == INT:
# Once we have the dense tensor, we can return it if we weren't given a
# fixed output sequence length. If we were, though, we have to dynamically
# choose whether to pad or trim it based on each tensor.
# We need to convert to dense if we have a ragged tensor.
if ragged_tensor.is_ragged(indexed_data):
dense_data = indexed_data.to_tensor(default_value=0)
else:
dense_data = indexed_data
if self._output_sequence_length is None:
return dense_data
else:
sequence_len = K.shape(dense_data)[1]
pad_amt = self._output_sequence_length - sequence_len
pad_fn = lambda: array_ops.pad(dense_data, [[0, 0], [0, pad_amt]])
slice_fn = lambda: dense_data[:, :self._output_sequence_length]
return control_flow_ops.cond(
sequence_len < self._output_sequence_length,
true_fn=pad_fn,
false_fn=slice_fn)
out_depth = self._max_tokens if self._pad_to_max else math_ops.cast(
(self._get_table_size() + self._reserved_values), dtypes.int32)
if self._output_mode == BINARY:
bool_one_hot_data = array_ops.one_hot(
indexed_data, depth=out_depth, on_value=True, off_value=False)
reduced_bool_data = math_ops.reduce_any(bool_one_hot_data, axis=1)
binary_data = math_ops.cast(reduced_bool_data, dtypes.int64)
return binary_data
one_hot_data = array_ops.one_hot(indexed_data, depth=out_depth)
counts = math_ops.reduce_sum(one_hot_data, axis=1)
if self._output_mode == COUNT:
return math_ops.cast(counts, dtypes.int64)
tf_idf_data = math_ops.multiply(counts, self._tf_idf_weights)
if self._output_mode == TFIDF:
return tf_idf_data
# We can only get here if we didn't recognize the passed mode.
raise ValueError("Unknown output mode %s" % self._output_mode)
def kl_divergence(distribution_a, distribution_b,
allow_nan_stats=True, name=None):
"""Get the KL-divergence KL(distribution_a || distribution_b).
If there is no KL method registered specifically for `type(distribution_a)`
and `type(distribution_b)`, then the class hierarchies of these types are
searched.
If one KL method is registered between any pairs of classes in these two
parent hierarchies, it is used.
If more than one such registered method exists, the method whose registered
classes have the shortest sum MRO paths to the input types is used.
If more than one such shortest path exists, the first method
identified in the search is used (favoring a shorter MRO distance to
`type(distribution_a)`).
Args:
distribution_a: The first distribution.
distribution_b: The second distribution.
allow_nan_stats: Python `bool`, default `True`. When `True`,
statistics (e.g., mean, mode, variance) use the value "`NaN`" to
indicate the result is undefined. When `False`, an exception is raised
if one or more of the statistic's batch members are undefined.
name: Python `str` name prefixed to Ops created by this class.
Returns:
A Tensor with the batchwise KL-divergence between `distribution_a`
and `distribution_b`.
Raises:
NotImplementedError: If no KL method is defined for distribution types
of `distribution_a` and `distribution_b`.
"""
kl_fn = _registered_kl(type(distribution_a), type(distribution_b))
if kl_fn is None:
raise NotImplementedError(
"No KL(distribution_a || distribution_b) registered for distribution_a "
"type %s and distribution_b type %s"
% (type(distribution_a).__name__, type(distribution_b).__name__))
with ops.name_scope("KullbackLeibler"):
kl_t = kl_fn(distribution_a, distribution_b, name=name)
if allow_nan_stats:
return kl_t
# Check KL for NaNs
kl_t = array_ops.identity(kl_t, name="kl")
with ops.control_dependencies([
control_flow_ops.Assert(
math_ops.logical_not(
math_ops.reduce_any(math_ops.is_nan(kl_t))),
["KL calculation between %s and %s returned NaN values "
"(and was called with allow_nan_stats=False). Values:"
% (distribution_a.name, distribution_b.name), kl_t])]):
return array_ops.identity(kl_t, name="checked_kl")
def next_inputs(self, time, outputs, state, sample_ids, name=None):
with ops.name_scope(name, "ScheduledOutputTrainingHelperNextInputs",
[time, outputs, state, sample_ids]):
(finished, base_next_inputs, state) = (
super(ScheduledOutputTrainingHelper, self).next_inputs(
time=time,
outputs=outputs,
state=state,
sample_ids=sample_ids,
name=name))
sample_ids = math_ops.cast(sample_ids, dtypes.bool)
def maybe_sample():
"""Perform scheduled sampling."""
def maybe_concatenate_auxiliary_inputs(outputs_, indices=None):
"""Concatenate outputs with auxiliary inputs, if they exist."""
if self._auxiliary_input_tas is None:
return outputs_
next_time = time + 1
auxiliary_inputs = nest.map_structure(
lambda ta: ta.read(next_time), self._auxiliary_input_tas)
if indices is not None:
auxiliary_inputs = array_ops.gather_nd(auxiliary_inputs, indices)
return nest.map_structure(
lambda x, y: array_ops.concat((x, y), -1),
outputs_, auxiliary_inputs)
if self._next_inputs_fn is None:
return array_ops.where(
sample_ids, maybe_concatenate_auxiliary_inputs(outputs),
base_next_inputs)
where_sampling = math_ops.cast(
array_ops.where(sample_ids), dtypes.int32)
where_not_sampling = math_ops.cast(
array_ops.where(math_ops.logical_not(sample_ids)), dtypes.int32)
outputs_sampling = array_ops.gather_nd(outputs, where_sampling)
inputs_not_sampling = array_ops.gather_nd(base_next_inputs,
where_not_sampling)
sampled_next_inputs = maybe_concatenate_auxiliary_inputs(
self._next_inputs_fn(outputs_sampling), where_sampling)
base_shape = array_ops.shape(base_next_inputs)
return (array_ops.scatter_nd(indices=where_sampling,
updates=sampled_next_inputs,
shape=base_shape)
+ array_ops.scatter_nd(indices=where_not_sampling,
updates=inputs_not_sampling,
shape=base_shape))
all_finished = math_ops.reduce_all(finished)
no_samples = math_ops.logical_not(math_ops.reduce_any(sample_ids))
next_inputs = control_flow_ops.cond(
math_ops.logical_or(all_finished, no_samples),
lambda: base_next_inputs, maybe_sample)
return (finished, next_inputs, state)
def kl(dist_a, dist_b, allow_nan=False, name=None):
"""Get the KL-divergence KL(dist_a || dist_b).
If there is no KL method registered specifically for `type(dist_a)` and
`type(dist_b)`, then the class hierarchies of these types are searched.
If one KL method is registered between any pairs of classes in these two
parent hierarchies, it is used.
If more than one such registered method exists, the method whose registered
classes have the shortest sum MRO paths to the input types is used.
If more than one such shortest path exists, the first method
identified in the search is used (favoring a shorter MRO distance to
`type(dist_a)`).
Args:
dist_a: The first distribution.
dist_b: The second distribution.
allow_nan: If `False` (default), a runtime error is raised
if the KL returns NaN values for any batch entry of the given
distributions. If `True`, the KL may return a NaN for the given entry.
name: (optional) Name scope to use for created operations.
Returns:
A Tensor with the batchwise KL-divergence between dist_a and dist_b.
Raises:
NotImplementedError: If no KL method is defined for distribution types
of dist_a and dist_b.
"""
kl_fn = _registered_kl(type(dist_a), type(dist_b))
if kl_fn is None:
raise NotImplementedError(
"No KL(dist_a || dist_b) registered for dist_a type %s and dist_b "
"type %s" % (type(dist_a).__name__, type(dist_b).__name__))
with ops.name_scope("KullbackLeibler"):
kl_t = kl_fn(dist_a, dist_b, name=name)
if allow_nan:
return kl_t
# Check KL for NaNs
kl_t = array_ops.identity(kl_t, name="kl")
with ops.control_dependencies([
control_flow_ops.Assert(
math_ops.logical_not(
math_ops.reduce_any(math_ops.is_nan(kl_t))),
[
"KL calculation between %s and %s returned NaN values "
"(and was called with allow_nan=False). Values:" %
(dist_a.name, dist_b.name), kl_t
])
]):
return array_ops.identity(kl_t, name="checked_kl")
def _detect_nan_inf(tensor):
"""Trace function for detecting any NaN/Inf in the tensor."""
if tensor.dtype.is_floating:
output_tensor = math_ops.reduce_any(
gen_math_ops.logical_or(gen_math_ops.is_nan(tensor),
gen_math_ops.is_inf(tensor)))
else:
output_tensor = constant_op.constant(False)
return _print_tensor(tensor_name, -1, tensor, output_tensor)
def kl(dist_a, dist_b, allow_nan=False, name=None):
"""Get the KL-divergence KL(dist_a || dist_b).
If there is no KL method registered specifically for `type(dist_a)` and
`type(dist_b)`, then the class hierarchies of these types are searched.
If one KL method is registered between any pairs of classes in these two
parent hierarchies, it is used.
If more than one such registered method exists, the method whose registered
classes have the shortest sum MRO paths to the input types is used.
If more than one such shortest path exists, the first method
identified in the search is used (favoring a shorter MRO distance to
`type(dist_a)`).
Args:
dist_a: The first distribution.
dist_b: The second distribution.
allow_nan: If `False` (default), a runtime error is raised
if the KL returns NaN values for any batch entry of the given
distributions. If `True`, the KL may return a NaN for the given entry.
name: (optional) Name scope to use for created operations.
Returns:
A Tensor with the batchwise KL-divergence between dist_a and dist_b.
Raises:
NotImplementedError: If no KL method is defined for distribution types
of dist_a and dist_b.
"""
kl_fn = _registered_kl(type(dist_a), type(dist_b))
if kl_fn is None:
raise NotImplementedError(
"No KL(dist_a || dist_b) registered for dist_a type %s and dist_b "
"type %s" % ((type(dist_a).__name__, type(dist_b).__name__)))
with ops.name_scope("KullbackLeibler"):
kl_t = kl_fn(dist_a, dist_b, name=name)
if allow_nan:
return kl_t
# Check KL for NaNs
kl_t = array_ops.identity(kl_t, name="kl")
with ops.control_dependencies([
control_flow_ops.Assert(
math_ops.logical_not(
math_ops.reduce_any(math_ops.is_nan(kl_t))),
["KL calculation between %s and %s returned NaN values "
"(and was called with allow_nan=False). Values:"
% (dist_a.name, dist_b.name), kl_t])]):
return array_ops.identity(kl_t, name="checked_kl")
def _detect_nan_inf(tensor):
"""Trace function for detecting any NaN/Inf in the tensor."""
if tensor.dtype.is_floating:
output_tensor = math_ops.reduce_any(
gen_math_ops.logical_or(
gen_math_ops.is_nan(tensor), gen_math_ops.is_inf(tensor)))
else:
output_tensor = constant_op.constant(False)
# The shape has to be 1. Set it if it does not have the information.
output_tensor = array_ops.reshape(output_tensor, [1])
return output_tensor
def _detect_nan_inf(tensor):
"""Trace function for detecting any NaN/Inf in the tensor."""
if tensor.dtype.is_floating:
# Since host can't handle bf16, always convert tensor to f32.
tensor = math_ops.cast(tensor, dtypes.float32)
output_tensor = math_ops.reduce_any(
gen_math_ops.logical_or(gen_math_ops.is_nan(tensor),
gen_math_ops.is_inf(tensor)))
else:
output_tensor = constant_op.constant(0)
return _print_tensor(op_name, output_idx, 1, tensor, output_tensor)
def _detect_nan_inf(tensor):
"""Trace function for detecting any NaN/Inf in the tensor."""
if tensor.dtype.is_floating:
# Since host can't handle bf16, always convert tensor to f32.
tensor = math_ops.cast(tensor, dtypes.float32)
output_tensor = math_ops.reduce_any(
gen_math_ops.logical_or(gen_math_ops.is_nan(tensor),
gen_math_ops.is_inf(tensor)))
else:
output_tensor = constant_op.constant(0)
return _print_tensor(op_name, output_idx, 1, tensor, output_tensor)
def _detect_nan_inf(tensor):
"""Trace function for detecting any NaN/Inf in the tensor."""
if tensor.dtype.is_floating:
output_tensor = math_ops.reduce_any(
gen_math_ops.logical_or(gen_math_ops.is_nan(tensor),
gen_math_ops.is_inf(tensor)))
else:
output_tensor = constant_op.constant(False)
# The shape has to be 1. Set it if it does not have the information.
output_tensor = array_ops.reshape(output_tensor, [1])
return output_tensor
def test_one_replica_eager_control_flow(self):
device = parallel_device.ParallelDevice(components=[
"/job:localhost/device:{}:0".format(self.device_type),
])
x = constant_op.constant([2, 3, 4])
with device:
x = device.pack([x])
if math_ops.reduce_any(math_ops.equal(x, constant_op.constant(4))):
y = constant_op.constant(1)
else:
y = constant_op.constant(2)
self.assertAllEqual([1], device.unpack(y))
def kl(dist_a, dist_b, allow_nan=False, name=None):
"""Get the KL-divergence KL(dist_a || dist_b).
Args:
dist_a: instance of distributions.Distribution.
dist_b: instance of distributions.Distribution.
allow_nan: If False (default), a runtime error is raised
if the KL returns NaN values for any batch entry of the given
distributions. If True, the KL may return a NaN for the given entry.
name: (optional) Name scope to use for created operations.
Returns:
A Tensor with the batchwise KL-divergence between dist_a and dist_b.
Raises:
TypeError: If dist_a or dist_b is not an instance of Distribution.
NotImplementedError: If no KL method is defined for distribution types
of dist_a and dist_b.
"""
if not isinstance(dist_a, distribution.Distribution):
raise TypeError(
"dist_a is not an instance of Distribution, received type: %s" %
type(dist_a))
if not isinstance(dist_b, distribution.Distribution):
raise TypeError(
"dist_b is not an instance of Distribution, received type: %s" %
type(dist_b))
kl_fn = _DIVERGENCES.get((type(dist_a), type(dist_b)), None)
if kl_fn is None:
raise NotImplementedError(
"No KL(dist_a || dist_b) registered for dist_a type %s and dist_b "
"type %s" % ((type(dist_a).__name__, type(dist_b).__name__)))
with ops.name_scope("KullbackLeibler"):
kl_t = kl_fn(dist_a, dist_b, name=name)
if allow_nan:
return kl_t
# Check KL for NaNs
kl_t = array_ops.identity(kl_t, name="kl")
with ops.control_dependencies([
logging_ops.Assert(
math_ops.logical_not(
math_ops.reduce_any(math_ops.is_nan(kl_t))),
[
"KL calculation between %s and %s returned NaN values "
"(and was called with allow_nan=False). Values:" %
(dist_a.name, dist_b.name), kl_t
])
]):
return array_ops.identity(kl_t, name="checked_kl")
def _compare(self, x, reduction_axes, keep_dims, use_gpu=False):
np_ans = x
if reduction_axes is None:
np_ans = np.any(np_ans, keepdims=keep_dims)
else:
for ra in reduction_axes[::-1]:
np_ans = np.any(np_ans, axis=ra, keepdims=keep_dims)
with self.test_session(use_gpu=use_gpu):
if reduction_axes is not None:
reduction_axes = np.array(reduction_axes).astype(np.int32)
tf_ans = math_ops.reduce_any(x, reduction_axes, keep_dims)
out = tf_ans.eval()
self.assertAllEqual(np_ans, out)
self.assertShapeEqual(np_ans, tf_ans)
def kl(dist_a, dist_b, allow_nan=False, name=None):
"""Get the KL-divergence KL(dist_a || dist_b).
Args:
dist_a: instance of distributions.Distribution.
dist_b: instance of distributions.Distribution.
allow_nan: If False (default), a runtime error is raised
if the KL returns NaN values for any batch entry of the given
distributions. If True, the KL may return a NaN for the given entry.
name: (optional) Name scope to use for created operations.
Returns:
A Tensor with the batchwise KL-divergence between dist_a and dist_b.
Raises:
TypeError: If dist_a or dist_b is not an instance of Distribution.
NotImplementedError: If no KL method is defined for distribution types
of dist_a and dist_b.
"""
if not isinstance(dist_a, distribution.Distribution):
raise TypeError("dist_a is not an instance of Distribution, received type: %s" % type(dist_a))
if not isinstance(dist_b, distribution.Distribution):
raise TypeError("dist_b is not an instance of Distribution, received type: %s" % type(dist_b))
kl_fn = _DIVERGENCES.get((type(dist_a), type(dist_b)), None)
if kl_fn is None:
raise NotImplementedError(
"No KL(dist_a || dist_b) registered for dist_a type %s and dist_b "
"type %s" % ((type(dist_a).__name__, type(dist_b).__name__))
)
with ops.name_scope("KullbackLeibler"):
kl_t = kl_fn(dist_a, dist_b, name=name)
if allow_nan:
return kl_t
# Check KL for NaNs
kl_t = array_ops.identity(kl_t, name="kl")
with ops.control_dependencies(
[
logging_ops.Assert(
math_ops.logical_not(math_ops.reduce_any(math_ops.is_nan(kl_t))),
[
"KL calculation between %s and %s returned NaN values "
"(and was called with allow_nan=False). Values:" % (dist_a.name, dist_b.name),
kl_t,
],
)
]
):
return array_ops.identity(kl_t, name="checked_kl")
def _compare(self, x, reduction_axes, keepdims, use_gpu=False):
np_ans = x
if reduction_axes is None:
np_ans = np.any(np_ans, keepdims=keepdims)
else:
for ra in reduction_axes[::-1]:
np_ans = np.any(np_ans, axis=ra, keepdims=keepdims)
with self.cached_session(use_gpu=use_gpu):
if reduction_axes is not None:
reduction_axes = np.array(reduction_axes).astype(np.int32)
tf_ans = math_ops.reduce_any(x, reduction_axes, keepdims)
out = self.evaluate(tf_ans)
self.assertAllEqual(np_ans, out)
self.assertShapeEqual(np_ans, tf_ans)
def next_inputs(self, time, outputs, state, sample_ids, name=None):
with ops.name_scope(name, "%sNextInputs" % type(self).__name__,
(time, outputs, state)):
def next_inputs(self, time, outputs, state):
next_time = time + 1
finished = (next_time >= self._sequence_length)
#all_finished = math_ops.reduce_all(finished)
def read_from_ta(inp):
return inp.read(next_time)
"""next_inputs = control_flow_ops.cond(
all_finished, lambda: self._zero_inputs,
lambda: nest.map_structure(read_from_ta, self._input_tas))"""
next_inputs = read_from_ta(self._input_tas)
return (finished, next_inputs, state)
(finished, base_next_inputs,
state) = next_inputs(self, time, outputs, state)
def maybe_sample():
where_sampling = math_ops.cast(array_ops.where(sample_ids),
dtypes.int32)
where_not_sampling = math_ops.cast(
array_ops.where(math_ops.logical_not(sample_ids)),
dtypes.int32)
outputs_sampling = array_ops.gather_nd(outputs, where_sampling)
inputs_not_sampling = array_ops.gather_nd(
base_next_inputs, where_not_sampling)
base_shape = array_ops.shape(base_next_inputs)
z_mean, z_logstd = tf.split(outputs_sampling, 2, 1)
z_logstd = tf.nn.elu(z_logstd)
sampled_next_inputs = z_mean + tf.random_normal(
tf.shape(z_mean)) * tf.exp(z_logstd)
return (array_ops.scatter_nd(indices=where_sampling,
updates=sampled_next_inputs,
shape=base_shape) +
array_ops.scatter_nd(indices=where_not_sampling,
updates=inputs_not_sampling,
shape=base_shape))
all_finished = math_ops.reduce_all(finished)
no_samples = math_ops.logical_not(math_ops.reduce_any(sample_ids))
next_inputs = control_flow_ops.cond(
math_ops.logical_or(all_finished, no_samples),
lambda: base_next_inputs, maybe_sample)
return (finished, next_inputs, state)
def convert_nan_or_inf_to_zero(self, grad):
"""Replace grad tensor with zero tensor if grad is NaN or Inf.
This is mainly for improving training stability. We skip updating the
variable by setting the grad to zero when there is NaN or Inf.
Args:
grad: Input gradient.
Returns:
a Tensor with the dtype equal to grad dtype.
"""
return array_ops.where(
math_ops.reduce_any(
math_ops.logical_or(math_ops.is_nan(grad),
math_ops.is_inf(grad))),
array_ops.zeros_like(grad, dtype=grad.dtype), grad)
def _process_scale(self, scale, event_ndims):
"""Helper to __init__ which gets scale in batch-ready form.
This function expands dimensions of `scale` according to the following
table:
event_ndims
scale.ndims 0 1
0 [1]+S+[1,1] "silent error"
1 [ ]+S+[1,1] "silent error"
2 [ ]+S+[1,1] [1]+S+[ ]
3 [ ]+S+[1,1] [ ]+S+[ ]
... (same) (same)
The idea is that we want to convert `scale` into something which can always
work for, say, the left-hand argument of `batch_matmul`.
Args:
scale: `Tensor`.
event_ndims: `Tensor` (0D, `int32`).
Returns:
scale: `Tensor` with dims expanded according to [above] table.
batch_ndims: `Tensor` (0D, `int32`). The ndims of the `batch` portion.
"""
ndims = array_ops.rank(scale)
left = math_ops.select(
math_ops.reduce_any([
math_ops.reduce_all([
math_ops.equal(ndims, 0),
math_ops.equal(event_ndims, 0)
]),
math_ops.reduce_all([
math_ops.equal(ndims, 2),
math_ops.equal(event_ndims, 1)
])]), 1, 0)
right = math_ops.select(math_ops.equal(event_ndims, 0), 2, 0)
pad = array_ops.concat(0, (
array_ops.ones([left], dtype=dtypes.int32),
array_ops.shape(scale),
array_ops.ones([right], dtype=dtypes.int32)))
scale = array_ops.reshape(scale, pad)
batch_ndims = ndims - 2 + right
return scale, batch_ndims
def _process_scale(self, scale, event_ndims):
"""Helper to __init__ which gets scale in batch-ready form.
This function expands dimensions of `scale` according to the following
table:
event_ndims
scale.ndims 0 1
0 [1]+S+[1,1] "silent error"
1 [ ]+S+[1,1] "silent error"
2 [ ]+S+[1,1] [1]+S+[ ]
3 [ ]+S+[1,1] [ ]+S+[ ]
... (same) (same)
The idea is that we want to convert `scale` into something which can always
work for, say, the left-hand argument of `batch_matmul`.
Args:
scale: `Tensor`.
event_ndims: `Tensor` (0D, `int32`).
Returns:
scale: `Tensor` with dims expanded according to [above] table.
batch_ndims: `Tensor` (0D, `int32`). The ndims of the `batch` portion.
"""
ndims = array_ops.rank(scale)
left = math_ops.select(
math_ops.reduce_any([
math_ops.reduce_all([
math_ops.equal(ndims, 0),
math_ops.equal(event_ndims, 0)
]),
math_ops.reduce_all([
math_ops.equal(ndims, 2),
math_ops.equal(event_ndims, 1)
])]), 1, 0)
right = math_ops.select(math_ops.equal(event_ndims, 0), 2, 0)
pad = array_ops.concat(0, (
array_ops.ones([left], dtype=dtypes.int32),
array_ops.shape(scale),
array_ops.ones([right], dtype=dtypes.int32)))
scale = array_ops.reshape(scale, pad)
batch_ndims = ndims - 2 + right
return scale, batch_ndims
def any(a, axis=None, keepdims=None): # pylint: disable=redefined-builtin
"""Whether any element in the entire array or in an axis evaluates to true.
Casts the array to bool type if it is not already and uses `tf.reduce_any` to
compute the result.
Args:
a: array_like. Could be an ndarray, a Tensor or any object that can be
converted to a Tensor using `tf.convert_to_tensor`.
axis: Optional. Could be an int or a tuple of integers. If not specified,
the reduction is performed over all array indices.
keepdims: If true, retains reduced dimensions with length 1.
Returns:
An ndarray. Note that unlike NumPy this does not return a scalar bool if
`axis` is None.
"""
a = asarray(a, dtype=bool)
return np_utils.tensor_to_ndarray(
math_ops.reduce_any(input_tensor=a.data, axis=axis, keepdims=keepdims))
def _assert_sparse_indices_are_ragged_right(indices):
"""Checks that the given SparseTensor.indices tensor is ragged-right.
Example: `indices = [[0, 0], [0, 1], [2, 0], [3, 1]]` is not ragged right
because the entry `[3, 1]` skips a cell.
Args:
indices: The SparseTensor indices to check.
Returns:
A list of control dependency op tensors.
"""
index_prefix = indices[:, :-1]
index_suffix = indices[:, -1]
# Check whether each index is starting a new row in the innermost dimension
# (prefix[i] != prefix[i-1]) or continuing a row (prefix[i] == prefix[i-1]).
# (Note: this skips the first index; we will check that separately below.)
index_prefix_changed = math_ops.reduce_any(math_ops.not_equal(
index_prefix[1:], index_prefix[:-1]),
axis=1)
# Check two cases:
# * For indices that start a new row: index_suffix[i] must be zero.
# * For indices that continue a row: index_suffix[i] must be equal to
# index_suffix[i-1]+1.
index_ok = array_ops.where(
index_prefix_changed, math_ops.equal(index_suffix[1:], 0),
math_ops.equal(index_suffix[1:], index_suffix[:-1] + 1))
# Also check that the very first index didn't skip any cells. The first
# index starts a new row (by definition), so its suffix should be zero.
sparse_indices_are_ragged_right = math_ops.logical_and(
math_ops.reduce_all(math_ops.equal(index_suffix[:1], 0)),
math_ops.reduce_all(index_ok))
message = [
'SparseTensor is not right-ragged', 'SparseTensor.indices =', indices
]
return [control_flow_ops.Assert(sparse_indices_are_ragged_right, message)]
def _check_weights_match_logits_and_reshape(weights, logits):
"""Checks that weights shape matches logits and reshapes if needed.
Consider logits of shape [D0, D1, ... DN, logits_dimension]. Weights shape
can be either:
* [D0, D1, ... DN, logits_dimension]
* [D0, D1, ... DN]: In this case, weights is reshaped into
[D0, D1, ... DN, 1] to work with weight broadcasting rules.
Args:
weights: weights Tensor.
logits: logits Tensor.
Returns:
Validated and reshaped weights Tensor.
"""
err_msg = (
'weights shape must be [D0, D1, ... DN], [D0, D1, ... DN, 1] or '
'[D0, D1, ... DN, logits_dimension]')
with ops.name_scope(None, 'weights', (weights, logits)) as scope:
weights_shape = array_ops.shape(weights, name='weights_shape')
logits_shape = array_ops.shape(logits, name='logits_shape')
if (weights.shape.ndims is not None and logits.shape.ndims is not None and
weights.shape.ndims == logits.shape.ndims - 1):
assert_dimension = check_ops.assert_equal(
logits_shape[:-1], weights_shape, message=err_msg,
data=['logits_shape: ', logits_shape,
'weights_shape: ', weights_shape])
with ops.control_dependencies([assert_dimension]):
return array_ops.expand_dims(weights, -1, name=scope)
supported_weights_shape = array_ops.concat([logits_shape[:-1], [1]], axis=0)
condition = math_ops.reduce_any(
[math_ops.reduce_all(math_ops.equal(logits_shape, weights_shape)),
math_ops.reduce_all(math_ops.equal(
supported_weights_shape, weights_shape))])
assert_dimension = control_flow_ops.Assert(
condition=condition,
data=[err_msg, 'logits_shape: ', logits_shape,
'weights_shape: ', weights_shape])
with ops.control_dependencies([assert_dimension]):
return array_ops.identity(weights, name=scope)
def _assert_sparse_indices_are_ragged_right(indices):
"""Checks that the given SparseTensor.indices tensor is ragged-right.
Example: `indices = [[0, 0], [0, 1], [2, 0], [3, 1]]` is not ragged right
because the entry `[3, 1]` skips a cell.
Args:
indices: The SparseTensor indices to check.
Returns:
A list of control dependency op tensors.
"""
index_prefix = indices[:, :-1]
index_suffix = indices[:, -1]
# Check whether each index is starting a new row in the innermost dimension
# (prefix[i] != prefix[i-1]) or continuing a row (prefix[i] == prefix[i-1]).
# (Note: this skips the first index; we will check that separately below.)
index_prefix_changed = math_ops.reduce_any(
math_ops.not_equal(index_prefix[1:], index_prefix[:-1]), axis=1)
# Check two cases:
# * For indices that start a new row: index_suffix[i] must be zero.
# * For indices that continue a row: index_suffix[i] must be equal to
# index_suffix[i-1]+1.
index_ok = array_ops.where(
index_prefix_changed, math_ops.equal(index_suffix[1:], 0),
math_ops.equal(index_suffix[1:], index_suffix[:-1] + 1))
# Also check that the very first index didn't skip any cells. The first
# index starts a new row (by definition), so its suffix should be zero.
sparse_indices_are_ragged_right = math_ops.logical_and(
math_ops.reduce_all(math_ops.equal(index_suffix[:1], 0)),
math_ops.reduce_all(index_ok))
message = [
'SparseTensor is not right-ragged',
'SparseTensor.indices =', indices
]
return [control_flow_ops.Assert(sparse_indices_are_ragged_right, message)]
def get_cluster_assignment(pairwise_distances, centroid_ids):
"""Assign data points to the neareset centroids.
Tensorflow has numerical instability and doesn't always choose
the data point with theoretically zero distance as it's nearest neighbor.
Thus, for each centroid in centroid_ids, explicitly assign
the centroid itself as the nearest centroid.
This is done through the mask tensor and the constraint_vect tensor.
Args:
pairwise_distances: 2-D Tensor of pairwise distances.
centroid_ids: 1-D Tensor of centroid indices.
Returns:
y_fixed: 1-D tensor of cluster assignment.
"""
predictions = math_ops.argmin(
array_ops.gather(pairwise_distances, centroid_ids), dimension=0)
batch_size = array_ops.shape(pairwise_distances)[0]
# Deal with numerical instability
mask = math_ops.reduce_any(array_ops.one_hot(
centroid_ids, batch_size, True, False, axis=-1, dtype=dtypes.bool),
axis=0)
constraint_one_hot = math_ops.multiply(
array_ops.one_hot(centroid_ids,
batch_size,
array_ops.constant(1, dtype=dtypes.int64),
array_ops.constant(0, dtype=dtypes.int64),
axis=0,
dtype=dtypes.int64),
math_ops.cast(math_ops.range(array_ops.shape(centroid_ids)[0]),
dtypes.int64))
constraint_vect = math_ops.reduce_sum(
array_ops.transpose(constraint_one_hot), axis=0)
y_fixed = array_ops.where(mask, constraint_vect, predictions)
return y_fixed
def insert(self, ids, scores):
"""Insert the ids and scores into the TopN."""
with ops.control_dependencies(self.last_ops):
scatter_op = state_ops.scatter_update(self.id_to_score, ids, scores)
larger_scores = math_ops.greater(scores, self.sl_scores[0])
def shortlist_insert():
larger_ids = array_ops.boolean_mask(
math_ops.to_int64(ids), larger_scores)
larger_score_values = array_ops.boolean_mask(scores, larger_scores)
shortlist_ids, new_ids, new_scores = tensor_forest_ops.top_n_insert(
self.sl_ids, self.sl_scores, larger_ids, larger_score_values)
u1 = state_ops.scatter_update(self.sl_ids, shortlist_ids, new_ids)
u2 = state_ops.scatter_update(self.sl_scores, shortlist_ids, new_scores)
return control_flow_ops.group(u1, u2)
# We only need to insert into the shortlist if there are any
# scores larger than the threshold.
cond_op = control_flow_ops.cond(
math_ops.reduce_any(larger_scores), shortlist_insert,
control_flow_ops.no_op)
with ops.control_dependencies([cond_op]):
self.last_ops = [scatter_op, cond_op]
def stopping_criterion(i, state):
return math_ops.logical_and(
i < max_iter,
math_ops.reduce_any(linalg.norm(state.r, axis=-1) > tol))
def _gcd_cond_fn(_, x2):
return math_ops.reduce_any(x2 != 0)
def get_next(self, name=None):
"""Returns the next input from the iterator for all replicas."""
if not self._enable_get_next_as_optional:
replicas = []
for i, worker in enumerate(self._input_workers.worker_devices):
if name is not None:
d = tf_device.DeviceSpec.from_string(worker)
new_name = "%s_%s_%d" % (name, d.job, d.task)
else:
new_name = None
with ops.device(worker):
# Make `replicas` a flat list of values across all replicas.
replicas.extend(
self._iterators[i].get_next_as_list_deprecated(new_name))
return values.regroup(self._input_workers.device_map, replicas)
replicas = []
worker_has_values = []
for i, worker in enumerate(self._input_workers.worker_devices):
if name is not None:
d = tf_device.DeviceSpec.from_string(worker)
new_name = "%s_%s_%d" % (name, d.job, d.task)
else:
new_name = None
with ops.device(worker):
worker_has_value, next_element = (
self._iterators[i].get_next_as_list(new_name))
worker_has_values.append(worker_has_value)
# Make `replicas` a flat list of values across all replicas.
replicas.append(next_element)
out_of_range_replicas = []
def out_of_range_fn(worker_index, device):
"""This function will throw an OutOfRange error."""
# As this will be only called when there is no data left, so calling
# get_next() will trigger an OutOfRange error.
data = self._iterators[worker_index].get_next(device)
out_of_range_replicas.append(data)
return data
# `global_has_value` indicates whether there is data in this global batch.
# We do a all-reduce across all the workers in the multi-worker case.
# TODO(b/126259107): Do strategy.reduce for CollectiveAllReduceStrategy.
if len(worker_has_values) > 1:
with ops.device(self._input_workers.compute_devices_for_worker(0)[0]):
# Place the tf.reduce_any op in device 0 to minimize communication
# cost.
# TODO(b/128545270): Investigate why placing it on worker 0 will cause
# the entire data to copy back from device to host.
global_has_value = math_ops.reduce_any(worker_has_values)
else:
global_has_value = worker_has_values[0]
results = []
for i, worker in enumerate(self._input_workers.worker_devices):
with ops.device(worker):
devices = self._input_workers.compute_devices_for_worker(i)
for j, device in enumerate(devices):
with ops.device(device):
# pylint: disable=undefined-loop-variable
# pylint: disable=cell-var-from-loop
# It is fine for the lambda to capture variables from the loop as
# the lambda is executed in the loop as well.
result = control_flow_ops.cond(global_has_value,
lambda: replicas[i][j],
lambda: out_of_range_fn(i, device))
# pylint: enable=cell-var-from-loop
# pylint: enable=undefined-loop-variable
results.append(result)
replicas = results
# Some dimensions in `replicas` will become unknown after we conditionally
# return the real tensors or the dummy tensors. We fix the input shapes by
# using the shapes from `out_of_range_replicas` because it is calling
# get_next() inside.
flattened_replicas = nest.flatten(replicas)
for i, replica_data in enumerate(nest.flatten(out_of_range_replicas)):
flattened_replicas[i].set_shape(replica_data.get_shape())
replicas = nest.pack_sequence_as(replicas, flattened_replicas)
return values.regroup(self._input_workers.device_map, replicas)
def _get_weights_and_check_match_logits(
features, weight_column, logits, allow_per_logit_weights=False):
"""Fetches weights from features and checks that the shape matches logits.
Consider logits of shape [D0, D1, ... DN, logits_dimension]. Weights shape
can be either:
* [D0, D1, ... DN, logits_dimension] if `allow_per_logit_weights=True`.
* [D0, D1, ... DN, 1]
* [D0, D1, ... DN]: In this case, weights is reshaped into
[D0, D1, ... DN, 1] to work with weight broadcasting rules.
Args:
features: The features dict that contains weights.
weight_column: The weight column. If not given, this method returns 1.
logits: logits Tensor.
allow_per_logit_weights: Boolean. Whether we allow weights along the logits
dimension, namely shape `[D0, D1, ... DN, logits_dimension]`.
Returns:
Validated and reshaped weights Tensor.
Raises:
ValueError: If the weights `Tensor` cannot be cast into float.
"""
if allow_per_logit_weights:
err_msg = (
'weights shape must be [D0, D1, ... DN], [D0, D1, ... DN, 1] or '
'[D0, D1, ... DN, logits_dimension]')
else:
err_msg = (
'weights shape must be [D0, D1, ... DN] or [D0, D1, ... DN, 1]')
with ops.name_scope(
None, 'weights',
values=tuple(six.itervalues(features)) + (logits,)) as scope:
# Fetch the weights.
if weight_column is None:
return 1.
if isinstance(weight_column, six.string_types):
weight_column = feature_column_lib.numeric_column(
key=weight_column, shape=(1,))
if not isinstance(weight_column, feature_column_lib._NumericColumn): # pylint: disable=protected-access
raise TypeError('Weight column must be either a string or _NumericColumn.'
' Given type: {}.'.format(type(weight_column)))
weights = weight_column._get_dense_tensor( # pylint: disable=protected-access
feature_column_lib._LazyBuilder(features)) # pylint: disable=protected-access
if not (weights.dtype.is_floating or weights.dtype.is_integer):
raise ValueError('Weight column should be castable to float. '
'Given dtype: {}'.format(weights.dtype))
weights = math_ops.to_float(weights, name='weights')
# Validate the weights shape.
weights_shape = array_ops.shape(weights, name='weights_shape')
logits_shape = array_ops.shape(logits, name='logits_shape')
if (weights.shape.ndims is not None and logits.shape.ndims is not None and
weights.shape.ndims == logits.shape.ndims - 1):
assert_dimension = check_ops.assert_equal(
logits_shape[:-1], weights_shape, message=err_msg,
data=['logits_shape: ', logits_shape,
'weights_shape: ', weights_shape])
with ops.control_dependencies([assert_dimension]):
return array_ops.expand_dims(weights, -1, name=scope)
supported_weights_shape = array_ops.concat([logits_shape[:-1], [1]], axis=0)
if allow_per_logit_weights:
condition = math_ops.reduce_any(
[math_ops.reduce_all(math_ops.equal(logits_shape, weights_shape)),
math_ops.reduce_all(math_ops.equal(
supported_weights_shape, weights_shape))])
assert_dimension = control_flow_ops.Assert(
condition=condition,
data=[err_msg, 'logits_shape: ', logits_shape,
'weights_shape: ', weights_shape])
else:
assert_dimension = check_ops.assert_equal(
supported_weights_shape, weights_shape, message=err_msg,
data=['logits_shape: ', logits_shape,
'weights_shape: ', weights_shape])
with ops.control_dependencies([assert_dimension]):
return array_ops.identity(weights, name=scope)
def get_next(self, name=None):
"""Returns the next input from the iterator for all replicas."""
if not self._enable_get_next_as_optional:
replicas = []
for i, worker in enumerate(self._input_workers.worker_devices):
if name is not None:
d = tf_device.DeviceSpec.from_string(worker)
new_name = "%s_%s_%d" % (name, d.job, d.task)
else:
new_name = None
with ops.device(worker):
# Make `replicas` a flat list of values across all replicas.
replicas.extend(
self._iterators[i].get_next_as_list_deprecated(new_name))
return values.regroup(self._input_workers.device_map, replicas)
replicas = []
worker_has_values = []
for i, worker in enumerate(self._input_workers.worker_devices):
if name is not None:
d = tf_device.DeviceSpec.from_string(worker)
new_name = "%s_%s_%d" % (name, d.job, d.task)
else:
new_name = None
with ops.device(worker):
worker_has_value, next_element = (
self._iterators[i].get_next_as_list(new_name))
worker_has_values.append(worker_has_value)
# Make `replicas` a flat list of values across all replicas.
replicas.append(next_element)
out_of_range_replicas = []
def out_of_range_fn(worker_index, device):
"""This function will throw an OutOfRange error."""
# As this will be only called when there is no data left, so calling
# get_next() will trigger an OutOfRange error.
data = self._iterators[worker_index].get_next(device)
out_of_range_replicas.append(data)
return data
# `global_has_value` indicates whether there is data in this global batch.
# We do a all-reduce across all the workers in the multi-worker case.
# TODO(b/126259107): Do strategy.reduce for CollectiveAllReduceStrategy.
if len(worker_has_values) > 1:
with ops.device(self._input_workers.compute_devices_for_worker(0)[0]):
# Place the tf.reduce_any op in device 0 to minimize communication
# cost.
# TODO(b/128545270): Investigate why placing it on worker 0 will cause
# the entire data to copy back from device to host.
global_has_value = math_ops.reduce_any(worker_has_values)
else:
global_has_value = worker_has_values[0]
results = []
for i, worker in enumerate(self._input_workers.worker_devices):
with ops.device(worker):
devices = self._input_workers.compute_devices_for_worker(i)
for j, device in enumerate(devices):
with ops.device(device):
# pylint: disable=undefined-loop-variable
# pylint: disable=cell-var-from-loop
# It is fine for the lambda to capture variables from the loop as
# the lambda is executed in the loop as well.
result = control_flow_ops.cond(global_has_value,
lambda: replicas[i][j],
lambda: out_of_range_fn(i, device))
# pylint: enable=cell-var-from-loop
# pylint: enable=undefined-loop-variable
results.append(result)
replicas = results
# Some dimensions in `replicas` will become unknown after we conditionally
# return the real tensors or the dummy tensors. We fix the input shapes by
# using the shapes from `out_of_range_replicas` because it is calling
# get_next() inside.
flattened_replicas = nest.flatten(replicas)
for i, replica_data in enumerate(nest.flatten(out_of_range_replicas)):
flattened_replicas[i].set_shape(replica_data.get_shape())
replicas = nest.pack_sequence_as(replicas, flattened_replicas)
return values.regroup(self._input_workers.device_map, replicas)
def terminate_when_all_zero(current_argument, residual_powers, accumulator):
del current_argument, accumulator # not used for condition
do_exit = math_ops.reduce_any(
math_ops.greater(residual_powers, array_ops.ones_like(residual_powers)))
return do_exit
def _cond(unused_chosen_inputs, running_products, unused_idx, unused_count): """Resampling loop exit condition.""" return math_ops.reduce_any(running_products > floor_vals)
def _initialized_limit_check():
return control_flow_ops.cond(
math_ops.reduce_any(non_decreasing),
lambda: state_ops.assign(discarded_windows_limiter, 0),
lambda: discarded_windows_limiter.count_up_to(self._discard_limit))
def build_controller(self):
"""RL optimization interface.
Returns:
ops: A dictionary holding handles of the model used for training.
"""
self._global_step = training_util.get_or_create_global_step()
ops = {}
ops["loss"] = 0
failing_signal = self.compute_reward(self.hparams.failing_signal)
ctr = {}
with tf_ops.name_scope("controller_{}".format(self.ctrl_id)):
with variable_scope.variable_scope("controller_{}".format(self.ctrl_id)):
ctr["reward"] = {"value": [], "ph": [], "update": []}
ctr["ready"] = {"value": [], "ph": [], "update": []}
ctr["best_reward"] = {"value": [], "update": []}
for i in range(self.hparams.num_children):
reward_value = variable_scope.get_local_variable(
"reward_{}".format(i),
initializer=0.0,
dtype=dtypes.float32,
trainable=False)
reward_ph = array_ops.placeholder(
dtypes.float32, shape=(), name="reward_ph_{}".format(i))
reward_update = state_ops.assign(
reward_value, reward_ph, use_locking=True)
ctr["reward"]["value"].append(reward_value)
ctr["reward"]["ph"].append(reward_ph)
ctr["reward"]["update"].append(reward_update)
best_reward = variable_scope.get_local_variable(
"best_reward_{}".format(i),
initializer=failing_signal,
dtype=dtypes.float32,
trainable=False)
ctr["best_reward"]["value"].append(best_reward)
ctr["best_reward"]["update"].append(
state_ops.assign(best_reward,
math_ops.minimum(best_reward, reward_update)))
ready_value = variable_scope.get_local_variable(
"ready_{}".format(i),
initializer=True,
dtype=dtypes.bool,
trainable=False)
ready_ph = array_ops.placeholder(
dtypes.bool, shape=(), name="ready_ph_{}".format(i))
ready_update = state_ops.assign(
ready_value, ready_ph, use_locking=True)
ctr["ready"]["value"].append(ready_value)
ctr["ready"]["ph"].append(ready_ph)
ctr["ready"]["update"].append(ready_update)
ctr["grouping_y_preds"], ctr["grouping_log_probs"] = self.get_groupings()
summary.histogram(
"grouping_actions",
array_ops.slice(ctr["grouping_y_preds"]["sample"], [0, 0],
[1, array_ops.shape(self.op_embeddings)[0]]))
with variable_scope.variable_scope("controller_{}".format(self.ctrl_id)):
ctr["baseline"] = variable_scope.get_local_variable(
"baseline",
initializer=failing_signal
if self.hparams.start_with_failing_signal else 0.0,
dtype=dtypes.float32,
trainable=False)
new_baseline = self.hparams.bl_dec * ctr["baseline"] + (
1 - self.hparams.bl_dec) * math_ops.reduce_mean(
ctr["reward"]["value"])
if not self.hparams.always_update_baseline:
baseline_mask = math_ops.less(ctr["reward"]["value"], failing_signal)
selected_reward = array_ops.boolean_mask(ctr["reward"]["value"],
baseline_mask)
selected_baseline = control_flow_ops.cond(
math_ops.reduce_any(baseline_mask),
lambda: math_ops.reduce_mean(selected_reward),
lambda: constant_op.constant(0, dtype=dtypes.float32))
ctr["pos_reward"] = selected_baseline
pos_ = math_ops.less(
constant_op.constant(0, dtype=dtypes.float32), selected_baseline)
selected_baseline = self.hparams.bl_dec * ctr["baseline"] + (
1 - self.hparams.bl_dec) * selected_baseline
selected_baseline = control_flow_ops.cond(
pos_, lambda: selected_baseline, lambda: ctr["baseline"])
new_baseline = control_flow_ops.cond(
math_ops.less(self.global_step,
self.hparams.stop_updating_after_steps),
lambda: new_baseline, lambda: selected_baseline)
ctr["baseline_update"] = state_ops.assign(
ctr["baseline"], new_baseline, use_locking=True)
ctr["y_preds"], ctr["log_probs"] = self.get_placements()
summary.histogram("actions", ctr["y_preds"]["sample"])
mask = math_ops.less(ctr["reward"]["value"], failing_signal)
ctr["loss"] = ctr["reward"]["value"] - ctr["baseline"]
ctr["loss"] *= (
ctr["log_probs"]["sample"] + ctr["grouping_log_probs"]["sample"])
selected_loss = array_ops.boolean_mask(ctr["loss"], mask)
selected_loss = control_flow_ops.cond(
math_ops.reduce_any(mask),
lambda: math_ops.reduce_mean(-selected_loss),
lambda: constant_op.constant(0, dtype=dtypes.float32))
ctr["loss"] = control_flow_ops.cond(
math_ops.less(self.global_step,
self.hparams.stop_updating_after_steps),
lambda: math_ops.reduce_mean(-ctr["loss"]), lambda: selected_loss)
ctr["reward_s"] = math_ops.reduce_mean(ctr["reward"]["value"])
summary.scalar("loss", ctr["loss"])
summary.scalar("avg_reward", ctr["reward_s"])
summary.scalar("best_reward_so_far", best_reward)
summary.scalar(
"advantage",
math_ops.reduce_mean(ctr["reward"]["value"] - ctr["baseline"]))
with variable_scope.variable_scope(
"optimizer", reuse=variable_scope.AUTO_REUSE):
(ctr["train_op"], ctr["lr"], ctr["grad_norm"],
ctr["grad_norms"]) = self._get_train_ops(
ctr["loss"],
tf_ops.get_collection(tf_ops.GraphKeys.TRAINABLE_VARIABLES),
self.global_step,
grad_bound=self.hparams.grad_bound,
lr_init=self.hparams.lr,
lr_dec=self.hparams.lr_dec,
start_decay_step=self.hparams.start_decay_step,
decay_steps=self.hparams.decay_steps,
optimizer_type=self.hparams.optimizer_type)
summary.scalar("gradnorm", ctr["grad_norm"])
summary.scalar("lr", ctr["lr"])
ctr["summary"] = summary.merge_all()
ops["controller"] = ctr
self.ops = ops
return ops
def test(self):
result_lt = ops.reduce_any(self.bool_lt, {'channel'})
golden_lt = core.LabeledTensor(
math_ops.reduce_any(self.bool_tensor, 1), [self.a0, self.a2, self.a3])
self.assertLabeledTensorsEqual(result_lt, golden_lt)
