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