def resize_and_center_cropped_inputs():
"""Deterministically resize to shorter side and center crop."""
input_shape = tf.compat.v1.shape(inputs)
input_height_t = input_shape[H_AXIS]
input_width_t = input_shape[W_AXIS]
ratio_cond = (input_height_t / input_width_t > (self.height / self.width))
# pylint: disable=g-long-lambda
resized_height = control_flow_util.smart_cond(
ratio_cond,
lambda: tf.cast(self.width * input_height_t / input_width_t,
input_height_t.dtype), lambda: self.height)
resized_width = control_flow_util.smart_cond(
ratio_cond, lambda: self.width,
lambda: tf.cast(self.height * input_width_t / input_height_t,
input_width_t.dtype))
# pylint: enable=g-long-lambda
resized_inputs = tf.image.resize(
images=inputs, size=tf.stack([resized_height, resized_width]))
img_hd_diff = resized_height - self.height
img_wd_diff = resized_width - self.width
bbox_h_start = tf.cast(img_hd_diff / 2, tf.int32)
bbox_w_start = tf.cast(img_wd_diff / 2, tf.int32)
bbox_begin = tf.stack([0, bbox_h_start, bbox_w_start, 0])
bbox_size = tf.stack([-1, self.height, self.width, -1])
outputs = tf.slice(resized_inputs, bbox_begin, bbox_size)
return outputs
def call(self, inputs, training=True):
if training is None:
training = backend.learning_phase()
def random_zoomed_inputs():
"""Zoomed inputs with random ops."""
inputs_shape = tf.compat.v1.shape(inputs)
batch_size = inputs_shape[0]
img_hd = tf.cast(inputs_shape[H_AXIS], tf.float32)
img_wd = tf.cast(inputs_shape[W_AXIS], tf.float32)
height_zoom = self._rng.uniform(
shape=[batch_size, 1],
minval=1. + self.height_lower,
maxval=1. + self.height_upper)
if self.width_factor is not None:
width_zoom = self._rng.uniform(
shape=[batch_size, 1],
minval=1. + self.width_lower,
maxval=1. + self.width_upper)
else:
width_zoom = height_zoom
zooms = tf.cast(
tf.concat([width_zoom, height_zoom], axis=1),
dtype=tf.float32)
return transform(
inputs,
get_zoom_matrix(zooms, img_hd, img_wd),
fill_mode=self.fill_mode,
fill_value=self.fill_value,
interpolation=self.interpolation)
output = control_flow_util.smart_cond(training, random_zoomed_inputs,
lambda: inputs)
output.set_shape(inputs.shape)
return output
def call(self, inputs, training=True):
if training is None:
training = backend.learning_phase()
def random_rotated_inputs():
"""Rotated inputs with random ops."""
inputs_shape = tf.compat.v1.shape(inputs)
batch_size = inputs_shape[0]
img_hd = tf.cast(inputs_shape[H_AXIS], tf.float32)
img_wd = tf.cast(inputs_shape[W_AXIS], tf.float32)
min_angle = self.lower * 2. * np.pi
max_angle = self.upper * 2. * np.pi
angles = self._rng.uniform(
shape=[batch_size], minval=min_angle, maxval=max_angle)
return transform(
inputs,
get_rotation_matrix(angles, img_hd, img_wd),
fill_mode=self.fill_mode,
fill_value=self.fill_value,
interpolation=self.interpolation)
output = control_flow_util.smart_cond(training, random_rotated_inputs,
lambda: inputs)
output.set_shape(inputs.shape)
return output
def wrap_with_training_arg(*args, **kwargs):
"""Wrap the `wrapped_call` function, and set training argument."""
try:
training = call_spec.get_arg_value("training",
args,
kwargs,
inputs_in_args=True)
except KeyError:
training = None
if training is None:
training = (default_training_value
or base_layer_utils.call_context().training
or backend.learning_phase())
args = list(args)
kwargs = kwargs.copy()
def replace_training_and_call(training):
new_args, new_kwargs = call_spec.set_arg_value("training",
training,
args,
kwargs,
inputs_in_args=True)
return wrapped_call(*new_args, **new_kwargs)
return control_flow_util.smart_cond(
training,
lambda: replace_training_and_call(True),
lambda: replace_training_and_call(False),
)
def _apply_gradients_cross_replica(self, distribution, grads_and_vars, name,
experimental_aggregate_gradients):
grads = [g for g, _ in grads_and_vars]
if isinstance(self._loss_scale, _DynamicLossScaleState):
loss_scale_update_op, should_apply_grads = self._loss_scale.update(grads)
else:
loss_scale_update_op = tf.no_op()
should_apply_grads = True
def apply_fn():
# We do not want DistributionStrategy to unwrap any MirroredVariables in
# grads_and_vars, because even in a replica context, the wrapped optimizer
# expects mirrored variables. So we wrap the variables with an
# _UnwrapPreventer, preventing DistributionStrategy from unwrapping the
# MirroredVariables.
wrapped_vars = _UnwrapPreventer([v for _, v in grads_and_vars])
return distribution.extended.call_for_each_replica(
self._apply_gradients,
args=(grads, wrapped_vars, name, experimental_aggregate_gradients))
def do_not_apply_fn():
# Normally self._optimizer.iterations is incremented in
# self._optimizer.apply_gradients(). Since that is not called in this
# branch, we increment it here instead.
return self._optimizer.iterations.assign_add(1, read_value=False)
# Note: We must call this cond() in a cross-replica context.
# DistributionStrategy does not support having a cond in a replica context
# with a branch that calls `merge_call`, and self._optimizer.apply_gradients
# calls `merge_call`.
maybe_apply_op = control_flow_util.smart_cond(
should_apply_grads, apply_fn, do_not_apply_fn)
return tf.group(maybe_apply_op, loss_scale_update_op)
def call(self, x, training=None):
if training is None:
training = keras.backend.learning_phase()
output = control_flow_util.smart_cond(training, lambda: x * 0,
lambda: tf.identity(x))
if not tf.executing_eagerly():
output._uses_learning_phase = True # pylint: disable=protected-access
return output
def call(self, inputs, training=True):
if training is None:
training = backend.learning_phase()
def random_cropped_inputs():
"""Cropped inputs with stateless random ops."""
input_shape = tf.compat.v1.shape(inputs)
crop_size = tf.stack(
[input_shape[0], self.height, self.width, input_shape[3]])
check = tf.Assert(
tf.reduce_all(input_shape >= crop_size),
[self.height, self.width])
with tf.control_dependencies([check]):
limit = input_shape - crop_size + 1
offset = stateless_random_ops.stateless_random_uniform(
tf.compat.v1.shape(input_shape),
dtype=crop_size.dtype,
maxval=crop_size.dtype.max,
seed=self._rng.make_seeds()[:, 0]) % limit
return tf.slice(inputs, offset, crop_size)
# TODO(b/143885775): Share logic with Resize and CenterCrop.
def resize_and_center_cropped_inputs():
"""Deterministically resize to shorter side and center crop."""
input_shape = tf.compat.v1.shape(inputs)
input_height_t = input_shape[H_AXIS]
input_width_t = input_shape[W_AXIS]
ratio_cond = (input_height_t / input_width_t > (self.height / self.width))
# pylint: disable=g-long-lambda
resized_height = control_flow_util.smart_cond(
ratio_cond,
lambda: tf.cast(self.width * input_height_t / input_width_t,
input_height_t.dtype), lambda: self.height)
resized_width = control_flow_util.smart_cond(
ratio_cond, lambda: self.width,
lambda: tf.cast(self.height * input_width_t / input_height_t,
input_width_t.dtype))
# pylint: enable=g-long-lambda
resized_inputs = tf.image.resize(
images=inputs, size=tf.stack([resized_height, resized_width]))
img_hd_diff = resized_height - self.height
img_wd_diff = resized_width - self.width
bbox_h_start = tf.cast(img_hd_diff / 2, tf.int32)
bbox_w_start = tf.cast(img_wd_diff / 2, tf.int32)
bbox_begin = tf.stack([0, bbox_h_start, bbox_w_start, 0])
bbox_size = tf.stack([-1, self.height, self.width, -1])
outputs = tf.slice(resized_inputs, bbox_begin, bbox_size)
return outputs
output = control_flow_util.smart_cond(training, random_cropped_inputs,
resize_and_center_cropped_inputs)
original_shape = inputs.shape.as_list()
batch_size, num_channels = original_shape[0], original_shape[3]
output_shape = [batch_size] + [self.height, self.width] + [num_channels]
output.set_shape(output_shape)
return output
def call(self, inputs, training=None, **kwargs):
if training is None:
training = backend.learning_phase()
outputs = smart_cond(training,
lambda: self._branch(inputs, self.train_scales),
lambda: self._branch(inputs, self.eval_scales))
return outputs
def _drop_attention_logits(self, logits: tf.Tensor, pad_mask: tf.Tensor,
training: tf.Tensor) -> tf.Tensor:
def droped_logits() -> tf.Tensor:
keep_prob = tf.random.uniform(tf.shape(logits), 0, 1) + pad_mask
drop_mask = tf.cast(
tf.less(keep_prob, self.attention_dropout_rate), logits.dtype)
return logits + drop_mask * -1e9
return smart_cond(training, droped_logits, lambda: tf.identity(logits))
def call(self, inputs, training=None, **kwargs):
if training is None:
training = tf.keras.backend.learning_phase()
x = smart_cond(
pred=training,
true_fn=lambda: self._transform(inputs),
false_fn=lambda: inputs,
)
return x
def call(self, inputs, training=None):
if training is None:
training = backend.learning_phase()
def dropped_inputs():
return self._random_generator.dropout(
inputs, self.rate, noise_shape=self._get_noise_shape(inputs))
output = control_flow_util.smart_cond(training, dropped_inputs,
lambda: tf.identity(inputs))
return output
def call(self, inputs, training=True):
if training is None:
training = backend.learning_phase()
def random_contrasted_inputs():
return tf.image.random_contrast(inputs, 1. - self.lower, 1. + self.upper,
self.seed)
output = control_flow_util.smart_cond(training, random_contrasted_inputs,
lambda: inputs)
output.set_shape(inputs.shape)
return output
def _apply_scores(self, scores, value, scores_mask=None, training=None):
"""Applies attention scores to the given value tensor.
To use this method in your attention layer, follow the steps:
* Use `query` tensor of shape `[batch_size, Tq]` and `key` tensor of
shape `[batch_size, Tv]` to calculate the attention `scores`.
* Pass `scores` and `value` tensors to this method. The method applies
`scores_mask`, calculates `attention_distribution = softmax(scores)`,
then returns `matmul(attention_distribution, value).
* Apply `query_mask` and return the result.
Args:
scores: Scores float tensor of shape `[batch_size, Tq, Tv]`.
value: Value tensor of shape `[batch_size, Tv, dim]`.
scores_mask: A boolean mask `Tensor` of shape `[batch_size, 1, Tv]` or
`[batch_size, Tq, Tv]`. If given, scores at positions where
`scores_mask==False` do not contribute to the result. It must
contain at least one `True` value in each line along the last
dimension.
training: Python boolean indicating whether the layer should behave in
training mode (adding dropout) or in inference mode (no dropout).
Returns:
Tensor of shape `[batch_size, Tq, dim]`.
Attention scores after masking and softmax with shape
`[batch_size, Tq, Tv]`.
"""
if scores_mask is not None:
padding_mask = tf.logical_not(scores_mask)
# Bias so padding positions do not contribute to attention
# distribution. Note 65504. is the max float16 value.
if scores.dtype is tf.float16:
scores -= 65504.0 * tf.cast(padding_mask, dtype=scores.dtype)
else:
scores -= 1.0e9 * tf.cast(padding_mask, dtype=scores.dtype)
if training is None:
training = backend.learning_phase()
weights = tf.nn.softmax(scores)
if self.dropout > 0:
def dropped_weights():
return self._random_generator.dropout(weights,
rate=self.dropout)
weights = control_flow_util.smart_cond(
training, dropped_weights, lambda: tf.identity(weights))
return tf.matmul(weights, value), weights
def call(self, inputs, *args, **kwargs):
if (1, 1) == self.output_size:
return self.case_global(inputs)
height, width = tf.unstack(tf.shape(inputs)[1:3])
pad_h = (self.output_size[0] - height % self.output_size[0]) % self.output_size[0]
pad_w = (self.output_size[1] - width % self.output_size[1]) % self.output_size[1]
# Hack to allow build non-divisible branch
inputs_ = tf.pad(inputs, [[0, 0], [0, pad_h], [0, pad_w], [0, 0]])
outputs = smart_cond(
(pad_h == 0) & (pad_w == 0),
lambda: self.case_divisible(inputs_),
lambda: self.case_nondivisible(inputs))
return outputs
def wrap_with_training_arg(*args, **kwargs):
"""Wrap the `wrapped_call` function, and set training argument."""
training_arg_index = get_training_arg_index(original_call)
training = get_training_arg(training_arg_index, args, kwargs)
if training is None:
training = default_training_value or backend.learning_phase()
args = list(args)
kwargs = kwargs.copy()
def replace_training_and_call(training):
set_training_arg(training, training_arg_index, args, kwargs)
return wrapped_call(*args, **kwargs)
return control_flow_util.smart_cond(
training, lambda: replace_training_and_call(True),
lambda: replace_training_and_call(False))
def call(self, inputs, training=True):
if training is None:
training = backend.learning_phase()
def random_flipped_inputs():
flipped_outputs = inputs
if self.horizontal:
flipped_outputs = tf.image.random_flip_left_right(
flipped_outputs, self.seed)
if self.vertical:
flipped_outputs = tf.image.random_flip_up_down(flipped_outputs,
self.seed)
return flipped_outputs
output = control_flow_util.smart_cond(training, random_flipped_inputs,
lambda: inputs)
output.set_shape(inputs.shape)
return output
def call(self, inputs, training=True):
if training is None:
training = backend.learning_phase()
def random_translated_inputs():
"""Translated inputs with random ops."""
inputs_shape = tf.compat.v1.shape(inputs)
batch_size = inputs_shape[0]
h_axis, w_axis = H_AXIS, W_AXIS
img_hd = tf.cast(inputs_shape[h_axis], tf.float32)
img_wd = tf.cast(inputs_shape[w_axis], tf.float32)
height_translate = self._rng.uniform(
shape=[batch_size, 1],
minval=self.height_lower,
maxval=self.height_upper,
dtype=tf.float32)
height_translate = height_translate * img_hd
width_translate = self._rng.uniform(
shape=[batch_size, 1],
minval=self.width_lower,
maxval=self.width_upper,
dtype=tf.float32)
width_translate = width_translate * img_wd
translations = tf.cast(
tf.concat([width_translate, height_translate], axis=1),
dtype=tf.float32)
return transform(
inputs,
get_translation_matrix(translations),
interpolation=self.interpolation,
fill_mode=self.fill_mode,
fill_value=self.fill_value)
output = control_flow_util.smart_cond(training, random_translated_inputs,
lambda: inputs)
output.set_shape(inputs.shape)
return output
def call(self, inputs, training=True):
if training is None:
training = backend.learning_phase()
def random_width_inputs():
"""Inputs width-adjusted with random ops."""
inputs_shape = tf.compat.v1.shape(inputs)
img_hd = inputs_shape[H_AXIS]
img_wd = tf.cast(inputs_shape[W_AXIS], tf.float32)
width_factor = self._rng.uniform(
shape=[],
minval=(1.0 + self.width_lower),
maxval=(1.0 + self.width_upper))
adjusted_width = tf.cast(width_factor * img_wd, tf.int32)
adjusted_size = tf.stack([img_hd, adjusted_width])
output = tf.image.resize(
images=inputs, size=adjusted_size, method=self._interpolation_method)
original_shape = inputs.shape.as_list()
output_shape = original_shape[0:2] + [None] + [original_shape[3]]
output.set_shape(output_shape)
return output
return control_flow_util.smart_cond(training, random_width_inputs,
lambda: inputs)
def call(self, inputs, training=True):
return control_flow_util.smart_cond(training, lambda: inputs * 0,
lambda: tf.identity(inputs))
def call(self, inputs, training=None, **kwargs):
if training is None:
training = backend.learning_phase()
image, mask, prev = inputs
no_prev = tf.logical_or(tf.cast(training, 'bool'),
tf.reduce_all(tf.equal(mask, prev)))
image = self._preprocess_image(image)
mask = self._preprocess_mask(mask)
"""
First iteration, s8 output
"""
pred_s8_1, scale_s8_1 = smart_cond(
no_prev, lambda: self._estimate_s8(image, mask),
lambda: self._preprocess_s8(prev))
"""
Second iteration, s4 output
"""
inp2 = tf.concat([image, mask, scale_s8_1, scale_s8_1], axis=-1)
_, feats4, feats8 = self.bone(inp2)
out2_s8 = self.psp(feats8)
mask_s8_2 = self.final8(out2_s8)
mask_s8_2 = resize_by_sample([mask_s8_2, image])
pred_s8_2 = self.act(mask_s8_2)
scale_s8_2 = self._preprocess_mask(
pred_s8_2 * 255.) # tanh in original implementation
out2_s4 = self.up1([out2_s8, feats4])
mask_s4_2 = self.final4(out2_s4)
mask_s4_2 = resize_by_sample([mask_s4_2, image])
pred_s4_2 = self.act(mask_s4_2)
scale_s4_2 = self._preprocess_mask(
pred_s4_2 * 255.) # tanh in original implementation
"""
Third iteration, s1 output
"""
inp3 = tf.concat([image, mask, scale_s8_2, scale_s4_2], axis=-1)
feats2, feats4, feats8 = self.bone(inp3)
out3_s8 = self.psp(feats8)
mask_s8_3 = self.final8(out3_s8)
mask_s8_3 = resize_by_sample([mask_s8_3, image])
pred_s8_3 = self.act(mask_s8_3)
out3_s4 = self.up1([out3_s8, feats4])
mask_s4_3 = self.final4(out3_s4)
mask_s4_3 = resize_by_sample([mask_s4_3, image])
pred_s4_3 = self.act(mask_s4_3)
out3_s4 = self.up2([out3_s4, feats2])
out3_s4 = self.up3([out3_s4, image])
mask_s1_3 = self.final1(tf.concat([out3_s4, image], axis=-1))
pred_s1_3 = self.act(mask_s1_3)
"""
Final output
"""
# Originally: pred_224 pred_56_2 pred_28_3 pred_56 pred_28_2 pred_28
preds = [
pred_s1_3, pred_s4_3, pred_s8_3, pred_s4_2, pred_s8_2, pred_s8_1
]
return preds
