jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def get_random_transformation_batch(
        self, batch_size, images=None, **kwargs
    ):
        heights, widths = augmentation_utils.get_images_shape(
            images, dtype=tf.float32
        )
        # get scaled_sizes
        if self.crop_to_aspect_ratio:
            scales = tf.maximum(self.height / heights, self.width / widths)
        elif self.pad_to_aspect_ratio:
            scales = tf.minimum(self.height / heights, self.width / widths)
        else:
            return {"dummy": tf.zeros((batch_size,))}
        new_heights = tf.cast(tf.round(heights * scales), tf.int32)
        new_widths = tf.cast(tf.round(widths * scales), tf.int32)
        scaled_sizes = tf.concat((new_heights, new_widths), axis=-1)

        # get padding values
        if self.crop_to_aspect_ratio:
            tops = tf.where(
                new_heights > self.height,
                tf.cast((new_heights - self.height) / 2, tf.int32),
                tf.constant(0, tf.int32),
            )
            bottoms = tf.where(
                new_heights > self.height,
                new_heights - self.height - tops,
                tf.constant(0, tf.int32),
            )
            lefts = tf.where(
                new_widths > self.width,
                tf.cast((new_widths - self.width) / 2, tf.int32),
                tf.constant(0, tf.int32),
            )
            rights = tf.where(
                new_widths > self.width,
                new_widths - self.width - lefts,
                tf.constant(0, tf.int32),
            )
        else:
            assert self.pad_to_aspect_ratio
            tops = tf.where(
                new_heights < self.height,
                tf.cast((self.height - new_heights) / 2, tf.int32),
                tf.constant(0, tf.int32),
            )
            bottoms = tf.where(
                new_heights < self.height,
                self.height - new_heights - tops,
                tf.constant(0, tf.int32),
            )
            lefts = tf.where(
                new_widths < self.width,
                tf.cast((self.width - new_widths) / 2, tf.int32),
                tf.constant(0, tf.int32),
            )
            rights = tf.where(
                new_widths < self.width,
                self.width - new_widths - lefts,
                tf.constant(0, tf.int32),
            )
        (tops, bottoms, lefts, rights) = augmentation_utils.get_position_params(
            tops, bottoms, lefts, rights, self.position, self._random_generator
        )
        return {
            "scaled_sizes": scaled_sizes,
            "tops": tops,
            "bottoms": bottoms,
            "lefts": lefts,
            "rights": rights,
        }
//  [python]
def mask_for_high_rank(val, val_mask, name=None):
    val_mask = tf.expand_dims(val_mask, -1)
    return tf.multiply(val, tf.to_float(val_mask), name=name or 'mask_for_high_rank')
//  [python]
def mask_for_high_rank(val, val_mask, name=None):
    val_mask = tf.expand_dims(val_mask, -1)
    return tf.multiply(val, tf.cast(val_mask, tf.float32), name=name or 'mask_for_high_rank')
//  [python]
def smooth_labels(label_smoothing, labels, question_lengths, token_ids):
  mask = 1. - (
      tf.to_float(tf.equal(token_ids, 0)) +
      tf.sequence_mask(question_lengths, labels.shape[-2], dtype=tf.float32))
  num_classes = tf.expand_dims(tf.math.reduce_sum(mask, -1, keepdims=True), -1)
  labels = (1. - label_smoothing) * labels + (label_smoothing / num_classes)
  return labels * tf.expand_dims(mask, -1)
//  [python]
def smooth_labels(label_smoothing, labels, question_lengths, token_ids):
  mask = 1. - (
      tf.cast(tf.equal(token_ids, 0), tf.float32) +
      tf.sequence_mask(question_lengths, labels.shape[-2], dtype=tf.float32))
  num_classes = tf.expand_dims(tf.math.reduce_sum(mask, -1, keepdims=True), -1)
  labels = (1. - label_smoothing) * labels + (label_smoothing / num_classes)
  return labels * tf.expand_dims(mask, -1)
//  [python]
def exp_mask_for_high_rank(val, val_mask, name=None):
    val_mask = tf.expand_dims(val_mask, -1)
    return tf.add(val, (1 - t.cast(val_mask, tf.float32)) * VERY_NEGATIVE_NUMBER, 
                name=name or 'exp_mask_for_high_rank')
//  [python]
def exp_mask_for_high_rank(val, val_mask, name=None):
    val_mask = tf.expand_dims(val_mask, -1)
    return tf.add(val, (1 - tf.cast(val_mask, tf.float32)) * VERY_NEGATIVE_NUMBER, 
                name=name or 'exp_mask_for_high_rank')
//  [python]
def exp_mask_for_high_rank_position(val, val_mask, position, name=None):
    val_mask = tf.expand_dims(val_mask, -1)
    position = tf.expand_dims(position, -1)
    return tf.add(tf.add(val, (1 - t.cast(val_mask, tf.float32)) * VERY_NEGATIVE_NUMBER), t.cast(position, tf.float32),
                name=name or 'exp_mask_for_high_rank_position')
//  [python]
def exp_mask_for_high_rank_position(val, val_mask, position, name=None):
    val_mask = tf.expand_dims(val_mask, -1)
    position = tf.expand_dims(position, -1)
    return tf.add(tf.add(val, (1 - tf.cast(val_mask, tf.float32)) * VERY_NEGATIVE_NUMBER), tf.cast(position, tf.float32),
                name=name or 'exp_mask_for_high_rank_position')
//  [python]
def get_losses(self, final_output, target, is_softmax=True):
        '''Returns the loss for a Siamese Network.
        Args:
            final_output: tensor that represent the final output of the image bundle.
            target: Tensor of target to be output by the siamese network.
            
        Returns:
            losses: list of tensors representing each loss component
        '''
        print('setting up losses...')
        self.target = target
        self.final_output = final_output
        with tf.variable_scope('losses'):
            if is_softmax:
                correct_prediction = tf.equal(tf.argmax(final_output,1), target)
                self.accuracy = tf.reduce_mean(tf.to_float(correct_prediction))

                siamese_loss = tf.reduce_mean(
                    tf.nn.sparse_softmax_cross_entropy_with_logits(
                                logits=final_output, 
                                labels=target,
                                name='softmax_loss'))
                self.siamese_loss = siamese_loss
            else:
                # If it's not softmax, it's l2 norm loss.
                self.accuracy = 0
                # self.l2_loss = tf.losses.mean_squared_error(
                    # final_output,
                    # target,
                    # scope='d1',
                    # loss_collection=tf.GraphKeys,
                    # reduction="none") 
                self.l2_loss = []
                for i in range(len(final_output)):
                    target[i] = tf.to_float(target[i])
                    final_output[i] = tf.to_float(final_output[i])
                    self.l2_loss.append( tf.norm(target[i] - final_output[i], axis=1) )
                #self.l2_loss_sum = tf.reduce_sum(self.l2_loss, 1)  

                
                siamese_loss = self.l2_loss
                self.robust_l2_loss = []
                if self.threshold is not None:
                    for i in range(len(siamese_loss)):

                        ind = tf.unstack(siamese_loss[i])
                        siamese_loss[i] = [ tf.cond(tf.greater(x, self.threshold), 
                                            lambda: self.threshold + self.threshold * tf.log(x / self.threshold),
                                            lambda: x) for x in ind ]
                        self.robust_l2_loss.append(siamese_loss[i])
                        siamese_loss[i] = tf.stack(siamese_loss[i])
                     
                self.siamese_losses = []
                for i in range(len(siamese_loss)):
                    self.siamese_losses.append( tf.reduce_sum(siamese_loss[i]) / self.cfg['batch_size'] )
                self.siamese_losses[1] = self.siamese_losses[1] * self.cycle_rate

        
        self.cycle_loss_total = tf.add_n(self.siamese_losses)
        tf.add_to_collection(tf.GraphKeys.LOSSES, self.cycle_loss_total)
        losses = [self.cycle_loss_total]
        return losses
//  [python]
def get_losses(self, final_output, target, is_softmax=True):
        '''Returns the loss for a Siamese Network.
        Args:
            final_output: tensor that represent the final output of the image bundle.
            target: Tensor of target to be output by the siamese network.
            
        Returns:
            losses: list of tensors representing each loss component
        '''
        print('setting up losses...')
        self.target = target
        self.final_output = final_output
        with tf.variable_scope('losses'):
            if is_softmax:
                correct_prediction = tf.equal(tf.argmax(final_output,1), target)
                self.accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

                siamese_loss = tf.reduce_mean(
                    tf.nn.sparse_softmax_cross_entropy_with_logits(
                                logits=final_output, 
                                labels=target,
                                name='softmax_loss'))
                self.siamese_loss = siamese_loss
            else:
                # If it's not softmax, it's l2 norm loss.
                self.accuracy = 0
                # self.l2_loss = tf.losses.mean_squared_error(
                    # final_output,
                    # target,
                    # scope='d1',
                    # loss_collection=tf.GraphKeys,
                    # reduction="none") 
                self.l2_loss = []
                for i in range(len(final_output)):
                    target[i] = tf.to_float(target[i])
                    final_output[i] = tf.to_float(final_output[i])
                    self.l2_loss.append( tf.norm(target[i] - final_output[i], axis=1) )
                #self.l2_loss_sum = tf.reduce_sum(self.l2_loss, 1)  

                
                siamese_loss = self.l2_loss
                self.robust_l2_loss = []
                if self.threshold is not None:
                    for i in range(len(siamese_loss)):

                        ind = tf.unstack(siamese_loss[i])
                        siamese_loss[i] = [ tf.cond(tf.greater(x, self.threshold), 
                                            lambda: self.threshold + self.threshold * tf.log(x / self.threshold),
                                            lambda: x) for x in ind ]
                        self.robust_l2_loss.append(siamese_loss[i])
                        siamese_loss[i] = tf.stack(siamese_loss[i])
                     
                self.siamese_losses = []
                for i in range(len(siamese_loss)):
                    self.siamese_losses.append( tf.reduce_sum(siamese_loss[i]) / self.cfg['batch_size'] )
                self.siamese_losses[1] = self.siamese_losses[1] * self.cycle_rate

        
        self.cycle_loss_total = tf.add_n(self.siamese_losses)
        tf.add_to_collection(tf.GraphKeys.LOSSES, self.cycle_loss_total)
        losses = [self.cycle_loss_total]
        return losses
//  [python]
def hf_compute_loss(self, labels, logits):
        loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(
            from_logits=True, reduction=tf.keras.losses.Reduction.NONE
        )
        if self.config.tf_legacy_loss:
            # make sure only labels that are not equal to -100 affect the loss
            active_loss = tf.not_equal(tf.reshape(labels, (-1,)), -100)
            reduced_logits = tf.boolean_mask(tf.reshape(logits, (-1, shape_list(logits)[2])), active_loss)
            labels = tf.boolean_mask(tf.reshape(labels, (-1,)), active_loss)
            return loss_fn(labels, reduced_logits)

        # Clip negative labels to zero here to avoid NaNs and errors - those positions will get masked later anyway
        unmasked_loss = loss_fn(tf.nn.relu(labels), logits)
        # make sure only labels that are not equal to -100 affect the loss
        loss_mask = tf.to_float(labels != -100)
        masked_loss = unmasked_loss * loss_mask
        reduced_masked_loss = tf.reduce_sum(masked_loss) / tf.reduce_sum(loss_mask)
        return tf.reshape(reduced_masked_loss, (1,))
//  [python]
def hf_compute_loss(self, labels, logits):
        loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(
            from_logits=True, reduction=tf.keras.losses.Reduction.NONE
        )
        if self.config.tf_legacy_loss:
            # make sure only labels that are not equal to -100 affect the loss
            active_loss = tf.not_equal(tf.reshape(labels, (-1,)), -100)
            reduced_logits = tf.boolean_mask(tf.reshape(logits, (-1, shape_list(logits)[2])), active_loss)
            labels = tf.boolean_mask(tf.reshape(labels, (-1,)), active_loss)
            return loss_fn(labels, reduced_logits)

        # Clip negative labels to zero here to avoid NaNs and errors - those positions will get masked later anyway
        unmasked_loss = loss_fn(tf.nn.relu(labels), logits)
        # make sure only labels that are not equal to -100 affect the loss
        loss_mask = tf.cast(labels != -100, dtype=unmasked_loss.dtype)
        masked_loss = unmasked_loss * loss_mask
        reduced_masked_loss = tf.reduce_sum(masked_loss) / tf.reduce_sum(loss_mask)
        return tf.reshape(reduced_masked_loss, (1,))
//  [python]
def test_int_support(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        for model_class in self.all_model_classes:
            prepared_for_class = self._prepare_for_class(
                inputs_dict.copy(),
                model_class,
                return_labels=True if "labels" in inspect.signature(model_class.call).parameters.keys() else False,
            )
            if not any(
                tensor.dtype.is_integer for tensor in prepared_for_class.values() if isinstance(tensor, tf.Tensor)
            ):
                return  # No integer inputs means no need for this test

            prepared_for_class = {
                key: tf.cast(tensor, tf.int64) if isinstance(tensor, tf.Tensor) and tensor.dtype.is_integer else tensor
                for key, tensor in prepared_for_class.items()
            }
            model = model_class(config)
            model(**prepared_for_class)  # No assertion, we're just checking this doesn't throw an error
            int32_prepared_for_class = {
                key: tf.cast(tensor, tf.int32) if isinstance(tensor, tf.Tensor) and tensor.dtype.is_integer else tensor
                for key, tensor in prepared_for_class.items()
            }
            model(**int32_prepared_for_class)  # No assertion, we're just checking this doesn't throw an error

            # After testing that the model accepts all int inputs, confirm that its dummies are int32
            for key, tensor in model.dummy_inputs.items():
                self.assertTrue(
                    isinstance(tensor, tf.Tensor) or keras.backend.is_keras_tensor(tensor),
                    "Dummy inputs should be tf.Tensor!",
                )
                if tensor.dtype.is_integer:
                    self.assertTrue(tensor.dtype == tf.int32, "Integer dummy inputs should be tf.int32!")

            # Also confirm that the input_signature uses int32
            for key, tensor_spec in model.input_signature.items():
                if tensor_spec.dtype.is_integer:
                    self.assertTrue(tensor_spec.dtype == tf.int32, "Input signatures should use tf.int32 for ints!")
//  [python]
def test_step(self, inputs):
        features, _, gt_boxes, gt_classes, *_ = inputs
        # tf.print(gt_boxes, gt_classes)
        # gt_boxes = tf.stack(
        #    [
        #        gt_boxes[..., 0] * self.config.image_size[0],
        #        gt_boxes[..., 1] * self.config.image_size[1],
        #        gt_boxes[..., 2] * self.config.image_size[0],
        #        gt_boxes[..., 3] * self.config.image_size[1],
        #    ],
        #    axis=-1,
        # )

        cls_out_list, box_out_list = self.call(features, training=False)
        ltrb, scores, classes, _ = postprocess.postprocess_per_class(
            self.config, cls_out_list, box_out_list
        )
        classes = tf.to_int32(classes)

        ap = tf.py_function(
            func=self._calc_mAP,
            inp=[ltrb, scores, classes, gt_boxes, gt_classes],
            Tout=tf.float64,
        )
        self.mAP_tracker.update_state(ap)

        return {m.name: m.result() for m in [self.mAP_tracker]}
//  [python]
def test_step(self, inputs):
        features, _, gt_boxes, gt_classes, *_ = inputs
        # tf.print(gt_boxes, gt_classes)
        # gt_boxes = tf.stack(
        #    [
        #        gt_boxes[..., 0] * self.config.image_size[0],
        #        gt_boxes[..., 1] * self.config.image_size[1],
        #        gt_boxes[..., 2] * self.config.image_size[0],
        #        gt_boxes[..., 3] * self.config.image_size[1],
        #    ],
        #    axis=-1,
        # )

        cls_out_list, box_out_list = self.call(features, training=False)
        ltrb, scores, classes, _ = postprocess.postprocess_per_class(
            self.config, cls_out_list, box_out_list
        )
        classes = tf.cast(classes, dtype=tf.int32)

        ap = tf.py_function(
            func=self._calc_mAP,
            inp=[ltrb, scores, classes, gt_boxes, gt_classes],
            Tout=tf.float64,
        )
        self.mAP_tracker.update_state(ap)

        return {m.name: m.result() for m in [self.mAP_tracker]}
//  [python]
def hf_compute_loss(self, labels, logits):
        loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(
            from_logits=True, reduction=tf.keras.losses.Reduction.NONE
        )
        if tf.executing_eagerly():  # Data-dependent conditionals are forbidden in XLA
            if tf.math.reduce_any(labels == -1):
                tf.print("Using `-1` to mask the loss for the token is deprecated. Please use `-100` instead.")

        if self.config.tf_legacy_loss:
            # make sure only labels that are not equal to -100
            # are taken into account as loss
            if tf.math.reduce_any(labels == -1):
                tf.print("Using `-1` to mask the loss for the token is deprecated. Please use `-100` instead.")
                active_loss = tf.reshape(labels, (-1,)) != -1
            else:
                active_loss = tf.reshape(labels, (-1,)) != -100
            reduced_logits = tf.boolean_mask(tf.reshape(logits, (-1, shape_list(logits)[2])), active_loss)
            labels = tf.boolean_mask(tf.reshape(labels, (-1,)), active_loss)

            return loss_fn(labels, reduced_logits)

        # Clip negative labels to zero here to avoid NaNs and errors - those positions will get masked later anyway
        unmasked_loss = loss_fn(tf.nn.relu(labels), logits)
        # make sure only labels that are not equal to -100 or -1
        # are taken into account as loss
        loss_mask = tf.to_float(labels >= 0)
        # Avoid possible division by zero later
        # Masked positions will have a loss of NaN because -100 and -1 are not valid labels
        masked_loss = unmasked_loss * loss_mask
        reduced_masked_loss = tf.reduce_sum(masked_loss) / tf.reduce_sum(loss_mask)
        return tf.reshape(reduced_masked_loss, (1,))
//  [python]
def hf_compute_loss(self, labels, logits):
        loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(
            from_logits=True, reduction=tf.keras.losses.Reduction.NONE
        )
        if tf.executing_eagerly():  # Data-dependent conditionals are forbidden in XLA
            if tf.math.reduce_any(labels == -1):
                tf.print("Using `-1` to mask the loss for the token is deprecated. Please use `-100` instead.")

        if self.config.tf_legacy_loss:
            # make sure only labels that are not equal to -100
            # are taken into account as loss
            if tf.math.reduce_any(labels == -1):
                tf.print("Using `-1` to mask the loss for the token is deprecated. Please use `-100` instead.")
                active_loss = tf.reshape(labels, (-1,)) != -1
            else:
                active_loss = tf.reshape(labels, (-1,)) != -100
            reduced_logits = tf.boolean_mask(tf.reshape(logits, (-1, shape_list(logits)[2])), active_loss)
            labels = tf.boolean_mask(tf.reshape(labels, (-1,)), active_loss)

            return loss_fn(labels, reduced_logits)

        # Clip negative labels to zero here to avoid NaNs and errors - those positions will get masked later anyway
        unmasked_loss = loss_fn(tf.nn.relu(labels), logits)
        # make sure only labels that are not equal to -100 or -1
        # are taken into account as loss
        loss_mask = tf.cast(labels >= 0, dtype=unmasked_loss.dtype)
        # Avoid possible division by zero later
        # Masked positions will have a loss of NaN because -100 and -1 are not valid labels
        masked_loss = unmasked_loss * loss_mask
        reduced_masked_loss = tf.reduce_sum(masked_loss) / tf.reduce_sum(loss_mask)
        return tf.reshape(reduced_masked_loss, (1,))
//  [python]
def hf_compute_loss(self, labels, logits):
        loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(
            from_logits=True, reduction=tf.keras.losses.Reduction.NONE
        )
        if self.config.tf_legacy_loss:
            # make sure only labels that are not equal to -100
            # are taken into account as loss
            next_sentence_active_loss = tf.not_equal(tf.reshape(labels, (-1,)), -100)
            next_sentence_reduced_logits = tf.boolean_mask(tf.reshape(logits, (-1, 2)), next_sentence_active_loss)
            next_sentence_label = tf.boolean_mask(tf.reshape(labels, (-1,)), next_sentence_active_loss)

            return loss_fn(next_sentence_label, next_sentence_reduced_logits)

        # make sure only labels that are not equal to -100
        # are taken into account as loss

        # Clip negative labels to zero here to avoid NaNs and errors - those positions will get masked later anyway
        unmasked_ns_loss = loss_fn(y_true=tf.nn.relu(labels), y_pred=logits)
        ns_loss_mask = tf.to_float(tf.cast(labels != -100, dtype=tf.bool))
        # Just zero out samples where label is -100, no reduction
        masked_ns_loss = unmasked_ns_loss * ns_loss_mask

        return masked_ns_loss
//  [python]
def hf_compute_loss(self, labels, logits):
        loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(
            from_logits=True, reduction=tf.keras.losses.Reduction.NONE
        )
        if self.config.tf_legacy_loss:
            # make sure only labels that are not equal to -100
            # are taken into account as loss
            next_sentence_active_loss = tf.not_equal(tf.reshape(labels, (-1,)), -100)
            next_sentence_reduced_logits = tf.boolean_mask(tf.reshape(logits, (-1, 2)), next_sentence_active_loss)
            next_sentence_label = tf.boolean_mask(tf.reshape(labels, (-1,)), next_sentence_active_loss)

            return loss_fn(next_sentence_label, next_sentence_reduced_logits)

        # make sure only labels that are not equal to -100
        # are taken into account as loss

        # Clip negative labels to zero here to avoid NaNs and errors - those positions will get masked later anyway
        unmasked_ns_loss = loss_fn(y_true=tf.nn.relu(labels), y_pred=logits)
        ns_loss_mask = tf.cast(labels != -100, dtype=unmasked_ns_loss.dtype)
        # Just zero out samples where label is -100, no reduction
        masked_ns_loss = unmasked_ns_loss * ns_loss_mask

        return masked_ns_loss
//  [python]
def _policy_loss(
      self, old_policy, policy, action, advantage, length):
    """Compute the policy loss composed of multiple components.

    1. The policy gradient loss is importance sampled from the data-collecting
       policy at the beginning of training.
    2. The second term is a KL penalty between the policy at the beginning of
       training and the current policy.
    3. Additionally, if this KL already changed more than twice the target
       amount, we activate a strong penalty discouraging further divergence.

    Args:
      old_policy: Action distribution of the behavioral policy.
      policy: Sequences of distribution params of the current policy.
      action: Sequences of actions.
      advantage: Sequences of advantages.
      length: Batch of sequence lengths.

    Returns:
      Tuple of loss tensor and summary tensor.
    """
    with tf.name_scope('policy_loss'):
      kl = tf.contrib.distributions.kl_divergence(old_policy, policy)
      # Infinite values in the KL, even for padding frames that we mask out,
      # cause NaN gradients since TensorFlow computes gradients with respect to
      # the whole input tensor.
      kl = tf.check_numerics(kl, 'kl')
      kl = tf.reduce_mean(self._mask(kl, length), 1)
      policy_gradient = tf.exp(
          policy.log_prob(action) - old_policy.log_prob(action))
      surrogate_loss = -tf.reduce_mean(self._mask(
          policy_gradient * tf.stop_gradient(advantage), length), 1)
      surrogate_loss = tf.check_numerics(surrogate_loss, 'surrogate_loss')
      kl_penalty = self._penalty * kl
      cutoff_threshold = self._config.kl_target * self._config.kl_cutoff_factor
      cutoff_count = tf.reduce_sum(
          tf.to_int32(kl > cutoff_threshold))
      with tf.control_dependencies([tf.cond(
          cutoff_count > 0,
          lambda: tf.Print(0, [cutoff_count], 'kl cutoff! '), int)]):
        kl_cutoff = (
            self._config.kl_cutoff_coef *
            tf.to_float(kl > cutoff_threshold) *
            (kl - cutoff_threshold) ** 2)
      policy_loss = surrogate_loss + kl_penalty + kl_cutoff
      entropy = tf.reduce_mean(policy.entropy(), axis=1)
      if self._config.entropy_regularization:
        policy_loss -= self._config.entropy_regularization * entropy
      summary = tf.summary.merge([
          tf.summary.histogram('entropy', entropy),
          tf.summary.histogram('kl', kl),
          tf.summary.histogram('surrogate_loss', surrogate_loss),
          tf.summary.histogram('kl_penalty', kl_penalty),
          tf.summary.histogram('kl_cutoff', kl_cutoff),
          tf.summary.histogram('kl_penalty_combined', kl_penalty + kl_cutoff),
          tf.summary.histogram('policy_loss', policy_loss),
          tf.summary.scalar('avg_surr_loss', tf.reduce_mean(surrogate_loss)),
          tf.summary.scalar('avg_kl_penalty', tf.reduce_mean(kl_penalty)),
          tf.summary.scalar('avg_policy_loss', tf.reduce_mean(policy_loss))])
      policy_loss = tf.reduce_mean(policy_loss, 0)
      return tf.check_numerics(policy_loss, 'policy_loss'), summary
//  [python]
def _policy_loss(
      self, old_policy, policy, action, advantage, length):
    """Compute the policy loss composed of multiple components.

    1. The policy gradient loss is importance sampled from the data-collecting
       policy at the beginning of training.
    2. The second term is a KL penalty between the policy at the beginning of
       training and the current policy.
    3. Additionally, if this KL already changed more than twice the target
       amount, we activate a strong penalty discouraging further divergence.

    Args:
      old_policy: Action distribution of the behavioral policy.
      policy: Sequences of distribution params of the current policy.
      action: Sequences of actions.
      advantage: Sequences of advantages.
      length: Batch of sequence lengths.

    Returns:
      Tuple of loss tensor and summary tensor.
    """
    with tf.name_scope('policy_loss'):
      kl = tf.contrib.distributions.kl_divergence(old_policy, policy)
      # Infinite values in the KL, even for padding frames that we mask out,
      # cause NaN gradients since TensorFlow computes gradients with respect to
      # the whole input tensor.
      kl = tf.check_numerics(kl, 'kl')
      kl = tf.reduce_mean(self._mask(kl, length), 1)
      policy_gradient = tf.exp(
          policy.log_prob(action) - old_policy.log_prob(action))
      surrogate_loss = -tf.reduce_mean(self._mask(
          policy_gradient * tf.stop_gradient(advantage), length), 1)
      surrogate_loss = tf.check_numerics(surrogate_loss, 'surrogate_loss')
      kl_penalty = self._penalty * kl
      cutoff_threshold = self._config.kl_target * self._config.kl_cutoff_factor
      cutoff_count = tf.reduce_sum(
          tf.cast(kl > cutoff_threshold, tf.int32))
      with tf.control_dependencies([tf.cond(
          cutoff_count > 0,
          lambda: tf.Print(0, [cutoff_count], 'kl cutoff! '), int)]):
        kl_cutoff = (
            self._config.kl_cutoff_coef *
            tf.cast(kl > cutoff_threshold, tf.float32) *
            (kl - cutoff_threshold) ** 2)
      policy_loss = surrogate_loss + kl_penalty + kl_cutoff
      entropy = tf.reduce_mean(policy.entropy(), axis=1)
      if self._config.entropy_regularization:
        policy_loss -= self._config.entropy_regularization * entropy
      summary = tf.summary.merge([
          tf.summary.histogram('entropy', entropy),
          tf.summary.histogram('kl', kl),
          tf.summary.histogram('surrogate_loss', surrogate_loss),
          tf.summary.histogram('kl_penalty', kl_penalty),
          tf.summary.histogram('kl_cutoff', kl_cutoff),
          tf.summary.histogram('kl_penalty_combined', kl_penalty + kl_cutoff),
          tf.summary.histogram('policy_loss', policy_loss),
          tf.summary.scalar('avg_surr_loss', tf.reduce_mean(surrogate_loss)),
          tf.summary.scalar('avg_kl_penalty', tf.reduce_mean(kl_penalty)),
          tf.summary.scalar('avg_policy_loss', tf.reduce_mean(policy_loss))])
      policy_loss = tf.reduce_mean(policy_loss, 0)
      return tf.check_numerics(policy_loss, 'policy_loss'), summary
//  [python]
def __build_meta_lstsq(self):
    """Build a meta least-square optimization problem."""

    # build a meta least-square optimization problem
    beta1 = 0.9
    beta2 = 0.999
    epsilon = 1e-8
    with tf.variable_scope('meta_lstsq'):
      # create placeholders to customize the least-square problem
      x_mat_ph = tf.placeholder(tf.float32, name='x_mat_ph')
      y_mat_ph = tf.placeholder(tf.float32, name='y_mat_ph')
      w_mat_ph = tf.placeholder(tf.float32, name='w_mat_ph')
      gacc1_ph = tf.placeholder(tf.float32, name='gacc1_ph')
      gacc2_ph = tf.placeholder(tf.float32, name='gacc2_ph')

      # create variables
      x_mat = tf.get_variable('x_mat', initializer=x_mat_ph, validate_shape=False)
      y_mat = tf.get_variable('y_mat', initializer=y_mat_ph, validate_shape=False)
      w_mat = tf.get_variable('w_mat', initializer=w_mat_ph, validate_shape=False)
      gacc1 = tf.get_variable('gacc1', initializer=gacc1_ph, validate_shape=False)
      gacc2 = tf.get_variable('gacc2', initializer=gacc2_ph, validate_shape=False)
      train_step = tf.get_variable('train_step', shape=[], initializer=tf.zeros_initializer)

      # TF operations
      nb_smpls = tf.to_float(tf.shape(x_mat)[0])
      loss_reg = tf.nn.l2_loss(tf.matmul(x_mat, w_mat) - y_mat) / nb_smpls
      loss_dcy = FLAGS.loss_w_dcy * tf.nn.l2_loss(w_mat)
      grad = tf.matmul(tf.transpose(x_mat), tf.matmul(x_mat, w_mat) - y_mat) / nb_smpls + FLAGS.loss_w_dcy * w_mat
      update_ops = [
        gacc1.assign(beta1 * gacc1 + (1.0 - beta1) * grad),
        gacc2.assign(beta2 * gacc2 + (1.0 - beta2) * grad ** 2),
        train_step.assign_add(tf.ones([]))
      ]
      with tf.control_dependencies(update_ops):
        lrn_rate = FLAGS.cpr_lstsq_lrn_rate \
          * tf.sqrt(1.0 - tf.pow(beta2, train_step)) / (1.0 - tf.pow(beta1, train_step))
        train_op = w_mat.assign_add(-lrn_rate * gacc1 / (tf.sqrt(gacc2) + epsilon))
      init_op = tf.variables_initializer([x_mat, y_mat, w_mat, gacc1, gacc2, train_step])

    # pack placeholders and variables into dict
    meta_lstsq = {
      'x_mat_ph': x_mat_ph,
      'y_mat_ph': y_mat_ph,
      'w_mat_ph': w_mat_ph,
      'gacc1_ph': gacc1_ph,
      'gacc2_ph': gacc2_ph,
      'w_mat': w_mat,
      'loss_reg': loss_reg,
      'loss_dcy': loss_dcy,
      'init_op': init_op,
      'train_op': train_op,
    }

    return meta_lstsq
//  [python]
def __build_meta_lstsq(self):
    """Build a meta least-square optimization problem."""

    # build a meta least-square optimization problem
    beta1 = 0.9
    beta2 = 0.999
    epsilon = 1e-8
    with tf.variable_scope('meta_lstsq'):
      # create placeholders to customize the least-square problem
      x_mat_ph = tf.placeholder(tf.float32, name='x_mat_ph')
      y_mat_ph = tf.placeholder(tf.float32, name='y_mat_ph')
      w_mat_ph = tf.placeholder(tf.float32, name='w_mat_ph')
      gacc1_ph = tf.placeholder(tf.float32, name='gacc1_ph')
      gacc2_ph = tf.placeholder(tf.float32, name='gacc2_ph')

      # create variables
      x_mat = tf.get_variable('x_mat', initializer=x_mat_ph, validate_shape=False)
      y_mat = tf.get_variable('y_mat', initializer=y_mat_ph, validate_shape=False)
      w_mat = tf.get_variable('w_mat', initializer=w_mat_ph, validate_shape=False)
      gacc1 = tf.get_variable('gacc1', initializer=gacc1_ph, validate_shape=False)
      gacc2 = tf.get_variable('gacc2', initializer=gacc2_ph, validate_shape=False)
      train_step = tf.get_variable('train_step', shape=[], initializer=tf.zeros_initializer)

      # TF operations
      nb_smpls = tf.cast(tf.shape(x_mat)[0], tf.float32)
      loss_reg = tf.nn.l2_loss(tf.matmul(x_mat, w_mat) - y_mat) / nb_smpls
      loss_dcy = FLAGS.loss_w_dcy * tf.nn.l2_loss(w_mat)
      grad = tf.matmul(tf.transpose(x_mat), tf.matmul(x_mat, w_mat) - y_mat) / nb_smpls + FLAGS.loss_w_dcy * w_mat
      update_ops = [
        gacc1.assign(beta1 * gacc1 + (1.0 - beta1) * grad),
        gacc2.assign(beta2 * gacc2 + (1.0 - beta2) * grad ** 2),
        train_step.assign_add(tf.ones([]))
      ]
      with tf.control_dependencies(update_ops):
        lrn_rate = FLAGS.cpr_lstsq_lrn_rate \
          * tf.sqrt(1.0 - tf.pow(beta2, train_step)) / (1.0 - tf.pow(beta1, train_step))
        train_op = w_mat.assign_add(-lrn_rate * gacc1 / (tf.sqrt(gacc2) + epsilon))
      init_op = tf.variables_initializer([x_mat, y_mat, w_mat, gacc1, gacc2, train_step])

    # pack placeholders and variables into dict
    meta_lstsq = {
      'x_mat_ph': x_mat_ph,
      'y_mat_ph': y_mat_ph,
      'w_mat_ph': w_mat_ph,
      'gacc1_ph': gacc1_ph,
      'gacc2_ph': gacc2_ph,
      'w_mat': w_mat,
      'loss_reg': loss_reg,
      'loss_dcy': loss_dcy,
      'init_op': init_op,
      'train_op': train_op,
    }

    return meta_lstsq
//  [python]
def denorm_boxes_graph(boxes, shape):
    """Converts boxes from normalized coordinates to pixel coordinates.
    boxes: [..., (y1, x1, y2, x2)] in normalized coordinates
    shape: [..., (height, width)] in pixels

    Note: In pixel coordinates (y2, x2) is outside the box. But in normalized
    coordinates it's inside the box.

    Returns:
        [..., (y1, x1, y2, x2)] in pixel coordinates
    """
    tensor_dtype = tf.float16 if K.floatx() != 'float32' else tf.float32

    h, w = tf.split(t.cast(shape, tensor_dtype), 2)
    scale = tf.concat([h, w, h, w], axis=-1) - tf.constant(1.0)
    shift = tf.constant([0., 0., 1., 1.])
    return t.cast(tf.round(tf.multiply(boxes, scale) + shift), tf.int32)
//  [python]
def denorm_boxes_graph(boxes, shape):
    """Converts boxes from normalized coordinates to pixel coordinates.
    boxes: [..., (y1, x1, y2, x2)] in normalized coordinates
    shape: [..., (height, width)] in pixels

    Note: In pixel coordinates (y2, x2) is outside the box. But in normalized
    coordinates it's inside the box.

    Returns:
        [..., (y1, x1, y2, x2)] in pixel coordinates
    """
    tensor_dtype = tf.float16 if K.floatx() != 'float32' else tf.float32

    h, w = tf.split(tf.cast(shape, tensor_dtype), 2)
    scale = tf.concat([h, w, h, w], axis=-1) - tf.constant(1.0)
    shift = tf.constant([0., 0., 1., 1.])
    return tf.cast(tf.round(tf.multiply(boxes, scale) + shift), tf.int32)
//  [python]
def loss_fun(logits, labels):
  """Add L2Loss to all the trainable variables.

  Add summary for "Loss" and "Loss/avg".
  Args:
    logits: Logits from inference().
    labels: Labels from distorted_inputs or inputs(). 1-D tensor
            of shape [batch_size]
    distillation: if set to True, use probabilities and not class labels to
                  compute softmax loss

  Returns:
    Loss tensor of type float.
  """

  # Calculate the cross entropy between labels and predictions
  labels = tf.to_int64(labels)
  cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(
      logits=logits, labels=labels, name='cross_entropy_per_example')

  # Calculate the average cross entropy loss across the batch.
  cross_entropy_mean = tf.reduce_mean(cross_entropy, name='cross_entropy')

  # Add to TF collection for losses
  tf.add_to_collection('losses', cross_entropy_mean)

  # The total loss is defined as the cross entropy loss plus all of the weight
  # decay terms (L2 loss).
  return tf.add_n(tf.get_collection('losses'), name='total_loss')
//  [python]
def loss_fun(logits, labels):
  """Add L2Loss to all the trainable variables.

  Add summary for "Loss" and "Loss/avg".
  Args:
    logits: Logits from inference().
    labels: Labels from distorted_inputs or inputs(). 1-D tensor
            of shape [batch_size]
    distillation: if set to True, use probabilities and not class labels to
                  compute softmax loss

  Returns:
    Loss tensor of type float.
  """

  # Calculate the cross entropy between labels and predictions
  labels = tf.cast(labels, tf.int64)
  cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(
      logits=logits, labels=labels, name='cross_entropy_per_example')

  # Calculate the average cross entropy loss across the batch.
  cross_entropy_mean = tf.reduce_mean(cross_entropy, name='cross_entropy')

  # Add to TF collection for losses
  tf.add_to_collection('losses', cross_entropy_mean)

  # The total loss is defined as the cross entropy loss plus all of the weight
  # decay terms (L2 loss).
  return tf.add_n(tf.get_collection('losses'), name='total_loss')
//  [python]
def mrcnn_one_mask_loss_graph(target_masks, target_class_ids, pred_masks):
    """Mask binary cross-entropy loss for the masks head.

    target_masks: [batch, num_rois, height, width].
        A float32 tensor of values 0 or 1. Uses zero padding to fill array.
    target_class_ids: [batch, num_rois]. Integer class IDs. Zero padded.
    pred_masks: [batch, proposals, height, width, num_classes] float32 tensor
                with values from 0 to 1.
    """
    # Reshape for simplicity. Merge first two dimensions into one.
    target_class_ids = K.reshape(target_class_ids, (-1,))
    target_class_ids = tf.clip_by_value(target_class_ids, clip_value_min=0, clip_value_max=1)
    mask_shape = tf.shape(target_masks)
    target_masks = K.reshape(target_masks, (-1, mask_shape[2], mask_shape[3]))
    pred_shape = tf.shape(pred_masks)
    pred_masks = K.reshape(pred_masks,
                           (-1, pred_shape[2], pred_shape[3], pred_shape[4]))
    # Permute predicted masks to [N, num_classes, height, width]
    pred_masks = tf.transpose(pred_masks, [0, 3, 1, 2])

    # Only positive ROIs contribute to the loss. And only
    # the class specific mask of each ROI.
    positive_ix = tf.where(target_class_ids > 0)[:, 0]
    positive_class_ids = tf.to_int64(
        tf.gather(target_class_ids, positive_ix))
    indices = tf.stack([positive_ix, positive_class_ids], axis=1)

    # Gather the masks (predicted and true) that contribute to loss
    y_true = tf.gather(target_masks, positive_ix)
    y_pred = tf.gather(pred_masks, indices)

    tensor_dtype = tf.float16 if K.floatx() != 'float32' else tf.float32

    # Compute binary cross entropy. If no positive ROIs, then return 0.
    # shape: [batch, roi, num_classes]
    loss = K.switch(tf.size(y_true) > 0,
                    K.binary_crossentropy(target=y_true, output=y_pred),
                    tf.constant(0.0, dtype=tensor_dtype))
    loss = K.mean(loss)
    return loss
//  [python]
def mrcnn_one_mask_loss_graph(target_masks, target_class_ids, pred_masks):
    """Mask binary cross-entropy loss for the masks head.

    target_masks: [batch, num_rois, height, width].
        A float32 tensor of values 0 or 1. Uses zero padding to fill array.
    target_class_ids: [batch, num_rois]. Integer class IDs. Zero padded.
    pred_masks: [batch, proposals, height, width, num_classes] float32 tensor
                with values from 0 to 1.
    """
    # Reshape for simplicity. Merge first two dimensions into one.
    target_class_ids = K.reshape(target_class_ids, (-1,))
    target_class_ids = tf.clip_by_value(target_class_ids, clip_value_min=0, clip_value_max=1)
    mask_shape = tf.shape(target_masks)
    target_masks = K.reshape(target_masks, (-1, mask_shape[2], mask_shape[3]))
    pred_shape = tf.shape(pred_masks)
    pred_masks = K.reshape(pred_masks,
                           (-1, pred_shape[2], pred_shape[3], pred_shape[4]))
    # Permute predicted masks to [N, num_classes, height, width]
    pred_masks = tf.transpose(pred_masks, [0, 3, 1, 2])

    # Only positive ROIs contribute to the loss. And only
    # the class specific mask of each ROI.
    positive_ix = tf.where(target_class_ids > 0)[:, 0]
    positive_class_ids = tf.cast(
        tf.gather(target_class_ids, positive_ix), tf.int64)
    indices = tf.stack([positive_ix, positive_class_ids], axis=1)

    # Gather the masks (predicted and true) that contribute to loss
    y_true = tf.gather(target_masks, positive_ix)
    y_pred = tf.gather(pred_masks, indices)

    tensor_dtype = tf.float16 if K.floatx() != 'float32' else tf.float32

    # Compute binary cross entropy. If no positive ROIs, then return 0.
    # shape: [batch, roi, num_classes]
    loss = K.switch(tf.size(y_true) > 0,
                    K.binary_crossentropy(target=y_true, output=y_pred),
                    tf.constant(0.0, dtype=tensor_dtype))
    loss = K.mean(loss)
    return loss
//  [python]
def compile(self, learning_rate, momentum):
        """Gets the model ready for training. Adds losses, regularization, and
        metrics. Then calls the Keras compile() function.
        """

        # Optimizer object
        if self.config.OPTIMIZER == 'SGD':
            optimizer = keras.optimizers.SGD(
                lr=learning_rate, momentum=momentum, clipnorm=self.config.GRADIENT_CLIP_NORM)
        else:
            optimizer = keras.optimizers.Adam(
                lr=learning_rate, amsgrad=True, clipnorm=self.config.GRADIENT_CLIP_NORM)

        # Add Losses
        # First, clear previously set losses to avoid duplication
        self.keras_model._losses = []
        self.keras_model._per_input_losses = {}
        loss_names = [
            "rpn_class_loss",  "rpn_bbox_loss",
            "mrcnn_class_loss", "mrcnn_bbox_loss", "mrcnn_mask_loss"]
        for name in loss_names:
            layer = self.keras_model.get_layer(name)
            if layer.output in self.keras_model.losses:
                continue
            loss = (
                tf.reduce_mean(layer.output, keepdims=True)
                * self.config.LOSS_WEIGHTS.get(name, 1.))
            self.keras_model.add_loss(loss)

        tensor_dtype = tf.float32
        # Add L2 Regularization
        # Skip gamma and beta weights of batch normalization layers.
        reg_losses = [
            keras.regularizers.l2(self.config.WEIGHT_DECAY)(w) / tf.to_float(tf.size(w))
            for w in self.keras_model.trainable_weights
            if 'gamma' not in w.name and 'beta' not in w.name]
        self.keras_model.add_loss(tf.add_n(reg_losses))

        # Compile
        self.keras_model.compile(
            optimizer=optimizer,
            loss=[None] * len(self.keras_model.outputs))

        # Add metrics for losses
        for name in loss_names:
            if name in self.keras_model.metrics_names:
                continue
            layer = self.keras_model.get_layer(name)
            self.keras_model.metrics_names.append(name)
            loss = (
                tf.reduce_mean(layer.output, keepdims=True)
                * self.config.LOSS_WEIGHTS.get(name, 1.))
            self.keras_model.metrics_tensors.append(loss)
//  [python]
def compile(self, learning_rate, momentum):
        """Gets the model ready for training. Adds losses, regularization, and
        metrics. Then calls the Keras compile() function.
        """

        # Optimizer object
        if self.config.OPTIMIZER == 'SGD':
            optimizer = keras.optimizers.SGD(
                lr=learning_rate, momentum=momentum, clipnorm=self.config.GRADIENT_CLIP_NORM)
        else:
            optimizer = keras.optimizers.Adam(
                lr=learning_rate, amsgrad=True, clipnorm=self.config.GRADIENT_CLIP_NORM)

        # Add Losses
        # First, clear previously set losses to avoid duplication
        self.keras_model._losses = []
        self.keras_model._per_input_losses = {}
        loss_names = [
            "rpn_class_loss",  "rpn_bbox_loss",
            "mrcnn_class_loss", "mrcnn_bbox_loss", "mrcnn_mask_loss"]
        for name in loss_names:
            layer = self.keras_model.get_layer(name)
            if layer.output in self.keras_model.losses:
                continue
            loss = (
                tf.reduce_mean(layer.output, keepdims=True)
                * self.config.LOSS_WEIGHTS.get(name, 1.))
            self.keras_model.add_loss(loss)

        tensor_dtype = tf.float32
        # Add L2 Regularization
        # Skip gamma and beta weights of batch normalization layers.
        reg_losses = [
            keras.regularizers.l2(self.config.WEIGHT_DECAY)(w) / tf.cast(tf.size(w), tensor_dtype)
            for w in self.keras_model.trainable_weights
            if 'gamma' not in w.name and 'beta' not in w.name]
        self.keras_model.add_loss(tf.add_n(reg_losses))

        # Compile
        self.keras_model.compile(
            optimizer=optimizer,
            loss=[None] * len(self.keras_model.outputs))

        # Add metrics for losses
        for name in loss_names:
            if name in self.keras_model.metrics_names:
                continue
            layer = self.keras_model.get_layer(name)
            self.keras_model.metrics_names.append(name)
            loss = (
                tf.reduce_mean(layer.output, keepdims=True)
                * self.config.LOSS_WEIGHTS.get(name, 1.))
            self.keras_model.metrics_tensors.append(loss)
//  [python]
def crf_sequence_score(inputs, tag_indices, sequence_lengths,
                       transition_params):
    """Computes the unnormalized score for a tag sequence.

    Args:
      inputs: A [batch_size, max_seq_len, num_tags] tensor of unary potentials
    to use as input to the CRF layer.
      tag_indices: A [batch_size, max_seq_len] matrix of tag indices for which
    we compute the unnormalized score.
      sequence_lengths: A [batch_size] vector of true sequence lengths.
      transition_params: 

    Returns:
      sequence_scores: A [batch_size] vector of unnormalized sequence scores.

    """
    tag_indices = tf.to_int32(tag_indices)
    sequence_lengths = tf.to_int32(sequence_lengths)

    # If max_seq_len is 1, we skip the score calculation and simply gather the
    # unary potentials of the single tag.
    def _single_seq_fn():
        batch_size = tf.shape(inputs, out_type=tag_indices.dtype)[0]

        example_inds = tf.reshape(
            tf.range(batch_size, dtype=tag_indices.dtype), [-1, 1])
        sequence_scores = tf.gather_nd(
            tf.squeeze(inputs, [1]),
            tf.concat([example_inds, tag_indices], axis=1))
        sequence_scores = tf.where(
            tf.less_equal(sequence_lengths, 0), tf.zeros_like(sequence_scores),
            sequence_scores)
        return sequence_scores

    def _multi_seq_fn():
        # Compute the scores of the given tag sequence.
        unary_scores = crf_unary_score(tag_indices, sequence_lengths, inputs)
        binary_scores = crf_binary_score(tag_indices, sequence_lengths,
                                         transition_params)
        sequence_scores = unary_scores + binary_scores
        return sequence_scores

    if inputs.shape[1] == 1:
        return _single_seq_fn()
    else:
        return _multi_seq_fn()