jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def autocontrast(image):
  """Implements Autocontrast function from PIL using TF ops.

  Args:
    image: A 3D uint8 tensor.
  Returns:
    The image after it has had autocontrast applied to it and will be of type
    uint8.
  """

  def scale_channel(image):
    """Scale the 2D image using the autocontrast rule."""
    # A possibly cheaper version can be done using cumsum/unique_with_counts
    # over the histogram values, rather than iterating over the entire image.
    # to compute mins and maxes.
    lo = tf.cast(tf.reduce_min(image), tf.float32)
    hi = tf.cast(tf.reduce_max(image), tf.float32)

    # Scale the image, making the lowest value 0 and the highest value 255.
    def scale_values(im):
      scale = 255.0 / (hi - lo)
      offset = -lo * scale
      im = tf.cast(im, tf.float32) * scale + offset
      im = tf.clip_by_value(im, 0.0, 255.0)
      return tf.cast(im, tf.uint8)

    result = tf.cond(hi > lo, lambda: scale_values(image), lambda: image)
    return result

  # Assumes RGB for now.  Scales each channel independently
  # and then stacks the result.
  s1 = scale_channel(image[:, :, 0])
  s2 = scale_channel(image[:, :, 1])
  s3 = scale_channel(image[:, :, 2])
  image = tf.stack([s1, s2, s3], 2)
  return image
//  [python]
def ctc_batch_cost(y_true, y_pred, input_length, label_length):
    """Runs CTC loss algorithm on each batch element.

    # Arguments
        y_true: tensor `(samples, max_string_length)`
            containing the truth labels.
        y_pred: tensor `(samples, time_steps, num_categories)`
            containing the prediction, or output of the softmax.
        input_length: tensor `(samples, 1)` containing the sequence length for
            each batch item in `y_pred`.
        label_length: tensor `(samples, 1)` containing the sequence length for
            each batch item in `y_true`.

    # Returns
        Tensor with shape (samples,1) containing the
            CTC loss of each element.
    """
    label_length = tf.to_int32(tf.squeeze(label_length))
    input_length = tf.to_int32(tf.squeeze(input_length))
    sparse_labels = tf.to_int32(ctc_label_dense_to_sparse(y_true, label_length))

    y_pred = tf.log(tf.transpose(y_pred, perm=[1, 0, 2]) + 1e-8)

    return tf.expand_dims(ctc.ctc_loss(inputs=y_pred,
                                       labels=sparse_labels,
                                       sequence_length=input_length), 1)
//  [python]
def ctc_batch_cost(y_true, y_pred, input_length, label_length):
    """Runs CTC loss algorithm on each batch element.

    # Arguments
        y_true: tensor `(samples, max_string_length)`
            containing the truth labels.
        y_pred: tensor `(samples, time_steps, num_categories)`
            containing the prediction, or output of the softmax.
        input_length: tensor `(samples, 1)` containing the sequence length for
            each batch item in `y_pred`.
        label_length: tensor `(samples, 1)` containing the sequence length for
            each batch item in `y_true`.

    # Returns
        Tensor with shape (samples,1) containing the
            CTC loss of each element.
    """
    label_length = tf.cast(tf.squeeze(label_length), tf.int32)
    input_length = tf.cast(tf.squeeze(input_length), tf.int32)
    sparse_labels = tf.cast(ctc_label_dense_to_sparse(y_true, label_length), tf.int32)

    y_pred = tf.log(tf.transpose(y_pred, perm=[1, 0, 2]) + 1e-8)

    return tf.expand_dims(ctc.ctc_loss(inputs=y_pred,
                                       labels=sparse_labels,
                                       sequence_length=input_length), 1)
//  [python]
def infer(self, features, *args, **kwargs):  # pylint: disable=arguments-differ
    """Produce predictions from the model by running it."""
    del args, kwargs
    # Inputs and features preparation needed to handle edge cases.
    if not features:
      features = {}
    hparams = self.hparams
    inputs_old = None
    if "inputs" in features and len(features["inputs"].shape) < 4:
      inputs_old = features["inputs"]
      features["inputs"] = tf.expand_dims(features["inputs"], 2)

    def logits_to_samples(logits, key):
      """Get samples from logits."""
      # If the last dimension is 1 then we're using L1/L2 loss.
      if common_layers.shape_list(logits)[-1] == 1:
        return tf.to_int32(tf.squeeze(logits, axis=-1))
      if key == "targets":
        return pixels_from_softmax(
            logits, gumbel_noise_factor=0.0,
            temperature=hparams.pixel_sampling_temperature)
      # Argmax in TF doesn't handle more than 5 dimensions yet.
      logits_shape = common_layers.shape_list(logits)
      argmax = tf.argmax(tf.reshape(logits, [-1, logits_shape[-1]]), axis=-1)
      return tf.reshape(argmax, logits_shape[:-1])

    # Get predictions.
    try:
      num_channels = hparams.problem.num_channels
    except AttributeError:
      num_channels = 1
    if "inputs" in features:
      inputs_shape = common_layers.shape_list(features["inputs"])
      targets_shape = [inputs_shape[0], hparams.video_num_target_frames,
                       inputs_shape[2], inputs_shape[3], num_channels]
    else:
      tf.logging.warn("Guessing targets shape as no inputs are given.")
      targets_shape = [hparams.batch_size,
                       hparams.video_num_target_frames, 1, 1, num_channels]

    features["targets"] = tf.zeros(targets_shape, dtype=tf.int32)
    reward_in_mod = "target_reward" in self.problem_hparams.modality
    action_in_mod = "target_action" in self.problem_hparams.modality
    if reward_in_mod:
      # TODO(lukaszkaiser): this is a hack. get the actual reward history.
      if "input_reward" not in features:
        features["input_reward"] = tf.zeros(
            [inputs_shape[0], inputs_shape[1], 1], dtype=tf.int32)
      features["target_reward"] = tf.zeros(
          [targets_shape[0], targets_shape[1], 1], dtype=tf.int32)
    if action_in_mod and "target_action" not in features:
      features["target_action"] = tf.zeros(
          [targets_shape[0], targets_shape[1], 1], dtype=tf.int32)
    logits, _ = self(features)  # pylint: disable=not-callable
    if isinstance(logits, dict):
      results = {}
      for k, v in six.iteritems(logits):
        results[k] = logits_to_samples(v, k)
        results["%s_logits" % k] = v
      # HACK: bypassing decoding issues.
      results["outputs"] = results["targets"]
      results["scores"] = results["targets"]
    else:
      results = logits_to_samples(logits, "targets")

    # Restore inputs to not confuse Estimator in edge cases.
    if inputs_old is not None:
      features["inputs"] = inputs_old

    # Return results.
    return results
//  [python]
def infer(self, features, *args, **kwargs):  # pylint: disable=arguments-differ
    """Produce predictions from the model by running it."""
    del args, kwargs
    # Inputs and features preparation needed to handle edge cases.
    if not features:
      features = {}
    hparams = self.hparams
    inputs_old = None
    if "inputs" in features and len(features["inputs"].shape) < 4:
      inputs_old = features["inputs"]
      features["inputs"] = tf.expand_dims(features["inputs"], 2)

    def logits_to_samples(logits, key):
      """Get samples from logits."""
      # If the last dimension is 1 then we're using L1/L2 loss.
      if common_layers.shape_list(logits)[-1] == 1:
        return tf.cast(tf.squeeze(logits, axis=-1), tf.int32)
      if key == "targets":
        return pixels_from_softmax(
            logits, gumbel_noise_factor=0.0,
            temperature=hparams.pixel_sampling_temperature)
      # Argmax in TF doesn't handle more than 5 dimensions yet.
      logits_shape = common_layers.shape_list(logits)
      argmax = tf.argmax(tf.reshape(logits, [-1, logits_shape[-1]]), axis=-1)
      return tf.reshape(argmax, logits_shape[:-1])

    # Get predictions.
    try:
      num_channels = hparams.problem.num_channels
    except AttributeError:
      num_channels = 1
    if "inputs" in features:
      inputs_shape = common_layers.shape_list(features["inputs"])
      targets_shape = [inputs_shape[0], hparams.video_num_target_frames,
                       inputs_shape[2], inputs_shape[3], num_channels]
    else:
      tf.logging.warn("Guessing targets shape as no inputs are given.")
      targets_shape = [hparams.batch_size,
                       hparams.video_num_target_frames, 1, 1, num_channels]

    features["targets"] = tf.zeros(targets_shape, dtype=tf.int32)
    reward_in_mod = "target_reward" in self.problem_hparams.modality
    action_in_mod = "target_action" in self.problem_hparams.modality
    if reward_in_mod:
      # TODO(lukaszkaiser): this is a hack. get the actual reward history.
      if "input_reward" not in features:
        features["input_reward"] = tf.zeros(
            [inputs_shape[0], inputs_shape[1], 1], dtype=tf.int32)
      features["target_reward"] = tf.zeros(
          [targets_shape[0], targets_shape[1], 1], dtype=tf.int32)
    if action_in_mod and "target_action" not in features:
      features["target_action"] = tf.zeros(
          [targets_shape[0], targets_shape[1], 1], dtype=tf.int32)
    logits, _ = self(features)  # pylint: disable=not-callable
    if isinstance(logits, dict):
      results = {}
      for k, v in six.iteritems(logits):
        results[k] = logits_to_samples(v, k)
        results["%s_logits" % k] = v
      # HACK: bypassing decoding issues.
      results["outputs"] = results["targets"]
      results["scores"] = results["targets"]
    else:
      results = logits_to_samples(logits, "targets")

    # Restore inputs to not confuse Estimator in edge cases.
    if inputs_old is not None:
      features["inputs"] = inputs_old

    # Return results.
    return results
//  [python]
def mean_iou(y_true, y_pred):
    prec = []
    for t in np.arange(0.5, 1.0, 0.05):
        y_pred_ = tf.to_int32(y_pred > t)
        score, up_opt = tf.metrics.mean_iou(y_true, y_pred_, 2)
        K.get_session().run(tf.local_variables_initializer())
        with tf.control_dependencies([up_opt]):
            score = tf.identity(score)
        prec.append(score)
    return K.mean(K.stack(prec), axis=0)
//  [python]
def mean_iou(y_true, y_pred):
    prec = []
    for t in np.arange(0.5, 1.0, 0.05):
        y_pred_ = tf.cast(y_pred > t, tf.int32)
        score, up_opt = tf.metrics.mean_iou(y_true, y_pred_, 2)
        K.get_session().run(tf.local_variables_initializer())
        with tf.control_dependencies([up_opt]):
            score = tf.identity(score)
        prec.append(score)
    return K.mean(K.stack(prec), axis=0)
//  [python]
def string_to_int_mapper(keys_to_map, mapping, num_oov_buckets=1, suffix="_id"):
  """Creates a mapping function to convert strings to ints in a tf.data.Dataset.

  For `dataset` outputs of type `str`, uses the list of strings in the given
  input `mapping` to look up the strings using tf.contrib.lookup and convert
  them to same-shape tensors of size tf.int32.

  Example:
    vocab = ['the', 'fox', 'jumped']
    dataset = dataset.map(string_to_int_mapper(['words'], mapping=vocab))
    dataset['words_id']  # <-- 'the' is mapped to 0, 'fox' to 1, etc...

  Args:
    keys_to_map: List of strings that are keys for tf.string Tensors to lookup.
    mapping: List of strings (or string tensors) to do the lookup. If the
        mapping is already a lookup table, then we directly use it.
    num_oov_buckets: Number of OOV buckets to use (default = 1).
    suffix: String to append to the given keys to indicate the mapped Tensors.

  Returns:
    _mapper: A mapping function that can be used with the tf.data.Dataset API.
  """
  if isinstance(mapping, LookupInterface):
    table = mapping
  else:
    table = contrib_lookup.index_table_from_tensor(
        mapping=mapping, num_oov_buckets=num_oov_buckets)

  def _mapper(dataset):
    for k in keys_to_map:
      dataset[k + suffix] = tf.to_int32(table.lookup(dataset[k]))
    return dataset
  return _mapper
//  [python]
def string_to_int_mapper(keys_to_map, mapping, num_oov_buckets=1, suffix="_id"):
  """Creates a mapping function to convert strings to ints in a tf.data.Dataset.

  For `dataset` outputs of type `str`, uses the list of strings in the given
  input `mapping` to look up the strings using tf.contrib.lookup and convert
  them to same-shape tensors of size tf.int32.

  Example:
    vocab = ['the', 'fox', 'jumped']
    dataset = dataset.map(string_to_int_mapper(['words'], mapping=vocab))
    dataset['words_id']  # <-- 'the' is mapped to 0, 'fox' to 1, etc...

  Args:
    keys_to_map: List of strings that are keys for tf.string Tensors to lookup.
    mapping: List of strings (or string tensors) to do the lookup. If the
        mapping is already a lookup table, then we directly use it.
    num_oov_buckets: Number of OOV buckets to use (default = 1).
    suffix: String to append to the given keys to indicate the mapped Tensors.

  Returns:
    _mapper: A mapping function that can be used with the tf.data.Dataset API.
  """
  if isinstance(mapping, LookupInterface):
    table = mapping
  else:
    table = contrib_lookup.index_table_from_tensor(
        mapping=mapping, num_oov_buckets=num_oov_buckets)

  def _mapper(dataset):
    for k in keys_to_map:
      dataset[k + suffix] = tf.cast(table.lookup(dataset[k]), tf.int32)
    return dataset
  return _mapper
//  [python]
def ctc_decode(y_pred, input_length, greedy=True, beam_width=100,
               top_paths=1):
    """Decodes the output of a softmax.

    Can use either greedy search (also known as best path)
    or a constrained dictionary search.

    # Arguments
        y_pred: tensor `(samples, time_steps, num_categories)`
            containing the prediction, or output of the softmax.
        input_length: tensor `(samples, )` containing the sequence length for
            each batch item in `y_pred`.
        greedy: perform much faster best-path search if `true`.
            This does not use a dictionary.
        beam_width: if `greedy` is `false`: a beam search decoder will be used
            with a beam of this width.
        top_paths: if `greedy` is `false`,
            how many of the most probable paths will be returned.

    # Returns
        Tuple:
            List: if `greedy` is `true`, returns a list of one element that
                contains the decoded sequence.
                If `false`, returns the `top_paths` most probable
                decoded sequences.
                Important: blank labels are returned as `-1`.
            Tensor `(top_paths, )` that contains
                the log probability of each decoded sequence.
    """
    y_pred = tf.log(tf.transpose(y_pred, perm=[1, 0, 2]) + 1e-8)
    input_length = tf.to_int32(input_length)

    if greedy:
        (decoded, log_prob) = ctc.ctc_greedy_decoder(
            inputs=y_pred,
            sequence_length=input_length)
    else:
        (decoded, log_prob) = ctc.ctc_beam_search_decoder(
            inputs=y_pred,
            sequence_length=input_length, beam_width=beam_width,
            top_paths=top_paths)

    decoded_dense = [tf.sparse_to_dense(st.indices, st.dense_shape, st.values, default_value=-1)
                     for st in decoded]
    return (decoded_dense, log_prob)
//  [python]
def ctc_decode(y_pred, input_length, greedy=True, beam_width=100,
               top_paths=1):
    """Decodes the output of a softmax.

    Can use either greedy search (also known as best path)
    or a constrained dictionary search.

    # Arguments
        y_pred: tensor `(samples, time_steps, num_categories)`
            containing the prediction, or output of the softmax.
        input_length: tensor `(samples, )` containing the sequence length for
            each batch item in `y_pred`.
        greedy: perform much faster best-path search if `true`.
            This does not use a dictionary.
        beam_width: if `greedy` is `false`: a beam search decoder will be used
            with a beam of this width.
        top_paths: if `greedy` is `false`,
            how many of the most probable paths will be returned.

    # Returns
        Tuple:
            List: if `greedy` is `true`, returns a list of one element that
                contains the decoded sequence.
                If `false`, returns the `top_paths` most probable
                decoded sequences.
                Important: blank labels are returned as `-1`.
            Tensor `(top_paths, )` that contains
                the log probability of each decoded sequence.
    """
    y_pred = tf.log(tf.transpose(y_pred, perm=[1, 0, 2]) + 1e-8)
    input_length = tf.cast(input_length, tf.int32)

    if greedy:
        (decoded, log_prob) = ctc.ctc_greedy_decoder(
            inputs=y_pred,
            sequence_length=input_length)
    else:
        (decoded, log_prob) = ctc.ctc_beam_search_decoder(
            inputs=y_pred,
            sequence_length=input_length, beam_width=beam_width,
            top_paths=top_paths)

    decoded_dense = [tf.sparse_to_dense(st.indices, st.dense_shape, st.values, default_value=-1)
                     for st in decoded]
    return (decoded_dense, log_prob)
//  [python]
def blend(image1, image2, factor):
  """Blend image1 and image2 using 'factor'.
  Factor can be above 0.0.  A value of 0.0 means only image1 is used.
  A value of 1.0 means only image2 is used.  A value between 0.0 and
  1.0 means we linearly interpolate the pixel values between the two
  images.  A value greater than 1.0 "extrapolates" the difference
  between the two pixel values, and we clip the results to values
  between 0 and 255.
  Args:
    image1: An image Tensor of type uint8.
    image2: An image Tensor of type uint8.
    factor: A floating point value above 0.0.
  Returns:
    A blended image Tensor of type uint8.
  """
  if factor == 0.0:
    return tf.convert_to_tensor(image1)
  if factor == 1.0:
    return tf.convert_to_tensor(image2)

  image1 = tf.to_float(image1)
  image2 = tf.to_float(image2)

  difference = image2 - image1
  scaled = factor * difference

  # Do addition in float.
  temp = tf.to_float(image1) + scaled

  # Interpolate
  if factor > 0.0 and factor < 1.0:
    # Interpolation means we always stay within 0 and 255.
    return tf.cast(temp, tf.uint8)

  # Extrapolate:
  #
  # We need to clip and then cast.
  return tf.cast(tf.clip_by_value(temp, 0.0, 255.0), tf.uint8)
//  [python]
def blend(image1, image2, factor):
  """Blend image1 and image2 using 'factor'.
  Factor can be above 0.0.  A value of 0.0 means only image1 is used.
  A value of 1.0 means only image2 is used.  A value between 0.0 and
  1.0 means we linearly interpolate the pixel values between the two
  images.  A value greater than 1.0 "extrapolates" the difference
  between the two pixel values, and we clip the results to values
  between 0 and 255.
  Args:
    image1: An image Tensor of type uint8.
    image2: An image Tensor of type uint8.
    factor: A floating point value above 0.0.
  Returns:
    A blended image Tensor of type uint8.
  """
  if factor == 0.0:
    return tf.convert_to_tensor(image1)
  if factor == 1.0:
    return tf.convert_to_tensor(image2)

  image1 = tf.cast(image1, tf.float32)
  image2 = tf.cast(image2, tf.float32)

  difference = image2 - image1
  scaled = factor * difference

  # Do addition in float.
  temp = tf.cast(image1, tf.float32) + scaled

  # Interpolate
  if factor > 0.0 and factor < 1.0:
    # Interpolation means we always stay within 0 and 255.
    return tf.cast(temp, tf.uint8)

  # Extrapolate:
  #
  # We need to clip and then cast.
  return tf.cast(tf.clip_by_value(temp, 0.0, 255.0), tf.uint8)
//  [python]
def autocontrast(image):
  """Implements Autocontrast function from PIL using TF ops.
  Args:
    image: A 3D uint8 tensor.
  Returns:
    The image after it has had autocontrast applied to it and will be of type
    uint8.
  """

  def scale_channel(image):
    """Scale the 2D image using the autocontrast rule."""
    # A possibly cheaper version can be done using cumsum/unique_with_counts
    # over the histogram values, rather than iterating over the entire image.
    # to compute mins and maxes.
    lo = tf.to_float(tf.reduce_min(image))
    hi = tf.to_float(tf.reduce_max(image))

    # Scale the image, making the lowest value 0 and the highest value 255.
    def scale_values(im):
      scale = 255.0 / (hi - lo)
      offset = -lo * scale
      im = tf.to_float(im) * scale + offset
      im = tf.clip_by_value(im, 0.0, 255.0)
      return tf.cast(im, tf.uint8)

    result = tf.cond(hi > lo, lambda: scale_values(image), lambda: image)
    return result

  # Assumes RGB for now.  Scales each channel independently
  # and then stacks the result.
  s1 = scale_channel(image[:, :, 0])
  s2 = scale_channel(image[:, :, 1])
  s3 = scale_channel(image[:, :, 2])
  image = tf.stack([s1, s2, s3], 2)
  return image
//  [python]
def autocontrast(image):
  """Implements Autocontrast function from PIL using TF ops.
  Args:
    image: A 3D uint8 tensor.
  Returns:
    The image after it has had autocontrast applied to it and will be of type
    uint8.
  """

  def scale_channel(image):
    """Scale the 2D image using the autocontrast rule."""
    # A possibly cheaper version can be done using cumsum/unique_with_counts
    # over the histogram values, rather than iterating over the entire image.
    # to compute mins and maxes.
    lo = tf.cast(tf.reduce_min(image), tf.float32)
    hi = tf.cast(tf.reduce_max(image), tf.float32)

    # Scale the image, making the lowest value 0 and the highest value 255.
    def scale_values(im):
      scale = 255.0 / (hi - lo)
      offset = -lo * scale
      im = tf.cast(im, tf.float32) * scale + offset
      im = tf.clip_by_value(im, 0.0, 255.0)
      return tf.cast(im, tf.uint8)

    result = tf.cond(hi > lo, lambda: scale_values(image), lambda: image)
    return result

  # Assumes RGB for now.  Scales each channel independently
  # and then stacks the result.
  s1 = scale_channel(image[:, :, 0])
  s2 = scale_channel(image[:, :, 1])
  s3 = scale_channel(image[:, :, 2])
  image = tf.stack([s1, s2, s3], 2)
  return image
//  [python]
def _decode_input_tensor_to_features_dict(feature_map, hparams):
  """Convert the interactive input format (see above) to a dictionary.

  Args:
    feature_map: dict with inputs.
    hparams: model hyperparameters

  Returns:
    A features dictionary, as expected by the decoder.
  """
  x = tf.convert_to_tensor(feature_map["inputs"])
  t = tf.convert_to_tensor(feature_map["target_tags"])

  # Add a third empty dimension.
  x = tf.to_int32(tf.expand_dims(x, axis=[2]))
  t = tf.to_int32(tf.expand_dims(t, axis=[2]))

  p_hparams = hparams.problem_hparams
  input_space_id = tf.constant(p_hparams.input_space_id)
  target_space_id = tf.constant(p_hparams.target_space_id)

  features = {}
  features["input_space_id"] = input_space_id
  features["target_space_id"] = target_space_id
  features["decode_length"] = tf.shape(x)[1] + 50
  features["inputs"] = x
  features["target_tags"] = t

  return features
//  [python]
def _decode_input_tensor_to_features_dict(feature_map, hparams):
  """Convert the interactive input format (see above) to a dictionary.

  Args:
    feature_map: dict with inputs.
    hparams: model hyperparameters

  Returns:
    A features dictionary, as expected by the decoder.
  """
  x = tf.convert_to_tensor(feature_map["inputs"])
  t = tf.convert_to_tensor(feature_map["target_tags"])

  # Add a third empty dimension.
  x = tf.cast(tf.expand_dims(x, axis=[2]), tf.int32)
  t = tf.cast(tf.expand_dims(t, axis=[2]), tf.int32)

  p_hparams = hparams.problem_hparams
  input_space_id = tf.constant(p_hparams.input_space_id)
  target_space_id = tf.constant(p_hparams.target_space_id)

  features = {}
  features["input_space_id"] = input_space_id
  features["target_space_id"] = target_space_id
  features["decode_length"] = tf.shape(x)[1] + 50
  features["inputs"] = x
  features["target_tags"] = t

  return features
//  [python]
def _build_train(self):
    print("-" * 80)
    print("Build train graph")
    logits = self._model(self.x_train, is_training=True)
    log_probs = tf.nn.sparse_softmax_cross_entropy_with_logits(
      logits=logits, labels=self.y_train)
    self.loss = tf.reduce_mean(log_probs)

    self.train_preds = tf.argmax(logits, axis=1)
    self.train_preds = tf.to_int32(self.train_preds)
    self.train_acc = tf.equal(self.train_preds, self.y_train)
    self.train_acc = tf.to_int32(self.train_acc)
    self.train_acc = tf.reduce_sum(self.train_acc)

    tf_variables = [var
        for var in tf.trainable_variables() if var.name.startswith(self.name)]
    self.num_vars = count_model_params(tf_variables)
    print("Model has {} params".format(self.num_vars))

    self.global_step = tf.Variable(
      0, dtype=tf.int32, trainable=False, name="global_step")
    self.train_op, self.lr, self.grad_norm, self.optimizer = get_train_ops(
      self.loss,
      tf_variables,
      self.global_step,
      clip_mode=self.clip_mode,
      grad_bound=self.grad_bound,
      l2_reg=self.l2_reg,
      lr_init=self.lr_init,
      lr_dec_start=self.lr_dec_start,
      lr_dec_every=self.lr_dec_every,
      lr_dec_rate=self.lr_dec_rate,
      lr_cosine=self.lr_cosine,
      lr_max=self.lr_max,
      lr_min=self.lr_min,
      lr_T_0=self.lr_T_0,
      lr_T_mul=self.lr_T_mul,
      num_train_batches=self.num_train_batches,
      optim_algo=self.optim_algo,
      sync_replicas=self.sync_replicas,
      num_aggregate=self.num_aggregate,
      num_replicas=self.num_replicas)
//  [python]
def _build_train(self):
    print("-" * 80)
    print("Build train graph")
    logits = self._model(self.x_train, is_training=True)
    log_probs = tf.nn.sparse_softmax_cross_entropy_with_logits(
      logits=logits, labels=self.y_train)
    self.loss = tf.reduce_mean(log_probs)

    self.train_preds = tf.argmax(logits, axis=1)
    self.train_preds = tf.cast(self.train_preds, tf.int32)
    self.train_acc = tf.equal(self.train_preds, self.y_train)
    self.train_acc = tf.cast(self.train_acc, tf.int32)
    self.train_acc = tf.reduce_sum(self.train_acc)

    tf_variables = [var
        for var in tf.trainable_variables() if var.name.startswith(self.name)]
    self.num_vars = count_model_params(tf_variables)
    print("Model has {} params".format(self.num_vars))

    self.global_step = tf.Variable(
      0, dtype=tf.int32, trainable=False, name="global_step")
    self.train_op, self.lr, self.grad_norm, self.optimizer = get_train_ops(
      self.loss,
      tf_variables,
      self.global_step,
      clip_mode=self.clip_mode,
      grad_bound=self.grad_bound,
      l2_reg=self.l2_reg,
      lr_init=self.lr_init,
      lr_dec_start=self.lr_dec_start,
      lr_dec_every=self.lr_dec_every,
      lr_dec_rate=self.lr_dec_rate,
      lr_cosine=self.lr_cosine,
      lr_max=self.lr_max,
      lr_min=self.lr_min,
      lr_T_0=self.lr_T_0,
      lr_T_mul=self.lr_T_mul,
      num_train_batches=self.num_train_batches,
      optim_algo=self.optim_algo,
      sync_replicas=self.sync_replicas,
      num_aggregate=self.num_aggregate,
      num_replicas=self.num_replicas)
//  [python]
def build_valid_rl(self, shuffle=False):
    print("-" * 80)
    print("Build valid graph on shuffled data")
    with tf.device("/cpu:0"):
      # shuffled valid data: for choosing validation model
      if not shuffle and self.data_format == "NCHW":
        self.images["valid_original"] = np.transpose(
          self.images["valid_original"], [0, 3, 1, 2])
      x_valid_shuffle, y_valid_shuffle = tf.train.shuffle_batch(
        [self.images["valid_original"], self.labels["valid_original"]],
        batch_size=self.batch_size,
        capacity=25000,
        enqueue_many=True,
        min_after_dequeue=0,
        num_threads=16,
        seed=self.seed,
        allow_smaller_final_batch=True,
      )

      def _pre_process(x):
        x = tf.pad(x, [[4, 4], [4, 4], [0, 0]])
        x = tf.random_crop(x, [32, 32, 3], seed=self.seed)
        x = tf.image.random_flip_left_right(x, seed=self.seed)
        if self.data_format == "NCHW":
          x = tf.transpose(x, [2, 0, 1])

        return x

      if shuffle:
        x_valid_shuffle = tf.map_fn(
          _pre_process, x_valid_shuffle, back_prop=False)

    logits = self._model(x_valid_shuffle, False, reuse=True)
    valid_shuffle_preds = tf.argmax(logits, axis=1)
    valid_shuffle_preds = tf.to_int32(valid_shuffle_preds)
    self.valid_shuffle_acc = tf.equal(valid_shuffle_preds, y_valid_shuffle)
    self.valid_shuffle_acc = tf.to_int32(self.valid_shuffle_acc)
    self.valid_shuffle_acc = tf.reduce_sum(self.valid_shuffle_acc)
//  [python]
def build_valid_rl(self, shuffle=False):
    print("-" * 80)
    print("Build valid graph on shuffled data")
    with tf.device("/cpu:0"):
      # shuffled valid data: for choosing validation model
      if not shuffle and self.data_format == "NCHW":
        self.images["valid_original"] = np.transpose(
          self.images["valid_original"], [0, 3, 1, 2])
      x_valid_shuffle, y_valid_shuffle = tf.train.shuffle_batch(
        [self.images["valid_original"], self.labels["valid_original"]],
        batch_size=self.batch_size,
        capacity=25000,
        enqueue_many=True,
        min_after_dequeue=0,
        num_threads=16,
        seed=self.seed,
        allow_smaller_final_batch=True,
      )

      def _pre_process(x):
        x = tf.pad(x, [[4, 4], [4, 4], [0, 0]])
        x = tf.random_crop(x, [32, 32, 3], seed=self.seed)
        x = tf.image.random_flip_left_right(x, seed=self.seed)
        if self.data_format == "NCHW":
          x = tf.transpose(x, [2, 0, 1])

        return x

      if shuffle:
        x_valid_shuffle = tf.map_fn(
          _pre_process, x_valid_shuffle, back_prop=False)

    logits = self._model(x_valid_shuffle, False, reuse=True)
    valid_shuffle_preds = tf.argmax(logits, axis=1)
    valid_shuffle_preds = tf.cast(valid_shuffle_preds, tf.int32)
    self.valid_shuffle_acc = tf.equal(valid_shuffle_preds, y_valid_shuffle)
    self.valid_shuffle_acc = tf.cast(self.valid_shuffle_acc, tf.int32)
    self.valid_shuffle_acc = tf.reduce_sum(self.valid_shuffle_acc)
//  [python]
def extract_relation_representations(input_layer, input_ids, tokenizer):
  """Extracts relation representation from sentence sequence layer."""
  entity_representations = []
  entity_marker_ids = tokenizer.convert_tokens_to_ids(["[E1]", "[E2]"])
  for entity_marker_id in entity_marker_ids:
    mask = tf.to_float(tf.equal(input_ids, entity_marker_id))
    mask = tf.broadcast_to(tf.expand_dims(mask, -1), tf.shape(input_layer))
    entity_representation = tf.reduce_max(
        mask * input_layer, axis=1, keepdims=True)
    entity_representations.append(entity_representation)

  output_layer = tf.concat(entity_representations, axis=2)
  output_layer = tf.squeeze(output_layer, [1])
  tf.logging.info("entity marker pooling AFTER output shape %s",
                  output_layer.shape)

  return output_layer
//  [python]
def extract_relation_representations(input_layer, input_ids, tokenizer):
  """Extracts relation representation from sentence sequence layer."""
  entity_representations = []
  entity_marker_ids = tokenizer.convert_tokens_to_ids(["[E1]", "[E2]"])
  for entity_marker_id in entity_marker_ids:
    mask = tf.cast(tf.equal(input_ids, entity_marker_id), tf.float32)
    mask = tf.broadcast_to(tf.expand_dims(mask, -1), tf.shape(input_layer))
    entity_representation = tf.reduce_max(
        mask * input_layer, axis=1, keepdims=True)
    entity_representations.append(entity_representation)

  output_layer = tf.concat(entity_representations, axis=2)
  output_layer = tf.squeeze(output_layer, [1])
  tf.logging.info("entity marker pooling AFTER output shape %s",
                  output_layer.shape)

  return output_layer
//  [python]
def file_based_input_fn_builder(input_file, seq_length, fewshot_num_classes,
                                fewshot_num_examples_per_class, drop_remainder):
  """Creates an `input_fn` closure to be passed to tf.Estimator."""

  # Add one for the 'query' example.
  fewshot_batch = fewshot_num_classes * fewshot_num_examples_per_class + 1
  name_to_features = {
      "input_ids": tf.FixedLenSequenceFeature([seq_length], tf.int64),
      "input_mask": tf.FixedLenSequenceFeature([seq_length], tf.int64),
      "segment_ids": tf.FixedLenSequenceFeature([seq_length], tf.int64),
      "guid": tf.FixedLenSequenceFeature([], tf.string),
  }

  def _decode_record(record, name_to_features):
    """Decodes a record to a TensorFlow example."""
    _, example = tf.parse_single_sequence_example(
        record, sequence_features=name_to_features)

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.to_int32(t)
      shape = tf.shape(example[name])
      # sequence_examples come with dynamic/unknown dimension which we reshape
      # to explicit dimension for the fewshot "batch" size.
      example[name] = tf.reshape(t, tf.concat([[fewshot_batch], shape[1:]], 0))

    return example

  def input_fn(params):
    """The actual input function."""
    d = tf.data.TFRecordDataset(input_file)
    d = d.apply(
        tf.data.experimental.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=params["batch_size"],
            drop_remainder=drop_remainder))

    return d

  return input_fn
//  [python]
def file_based_input_fn_builder(input_file, seq_length, fewshot_num_classes,
                                fewshot_num_examples_per_class, drop_remainder):
  """Creates an `input_fn` closure to be passed to tf.Estimator."""

  # Add one for the 'query' example.
  fewshot_batch = fewshot_num_classes * fewshot_num_examples_per_class + 1
  name_to_features = {
      "input_ids": tf.FixedLenSequenceFeature([seq_length], tf.int64),
      "input_mask": tf.FixedLenSequenceFeature([seq_length], tf.int64),
      "segment_ids": tf.FixedLenSequenceFeature([seq_length], tf.int64),
      "guid": tf.FixedLenSequenceFeature([], tf.string),
  }

  def _decode_record(record, name_to_features):
    """Decodes a record to a TensorFlow example."""
    _, example = tf.parse_single_sequence_example(
        record, sequence_features=name_to_features)

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.cast(t, tf.int32)
      shape = tf.shape(example[name])
      # sequence_examples come with dynamic/unknown dimension which we reshape
      # to explicit dimension for the fewshot "batch" size.
      example[name] = tf.reshape(t, tf.concat([[fewshot_batch], shape[1:]], 0))

    return example

  def input_fn(params):
    """The actual input function."""
    d = tf.data.TFRecordDataset(input_file)
    d = d.apply(
        tf.data.experimental.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=params["batch_size"],
            drop_remainder=drop_remainder))

    return d

  return input_fn
//  [python]
def preprocess_for_test_v0(image, shape, output_height, output_width, data_format='NCHW', bbox_border=25., heatmap_sigma=1., heatmap_size=64, scope=None):
  """Preprocesses the given image for evaluation.

  Args:
    image: A `Tensor` representing an image of arbitrary size.
    output_height: The height of the image after preprocessing.
    output_width: The width of the image after preprocessing.

  Returns:
    A preprocessed image.
  """
  with tf.name_scope(scope, 'vgg_test_image', [image, output_height, output_width]):
    # Crop the central region of the image with an area containing 87.5% of
    # the original image.

    image = tf.expand_dims(image, 0)
    image = tf.image.resize_bilinear(image, [output_height, output_width], align_corners=False)
    image = tf.squeeze(image, [0])
    image.set_shape([output_height, output_width, 3])
    image = tf.to_float(image)

    normarlized_image = _mean_image_subtraction(image, [_R_MEAN, _G_MEAN, _B_MEAN])
    if data_format == 'NCHW':
      normarlized_image = tf.transpose(normarlized_image, perm=(2, 0, 1))
    return normarlized_image/255.
//  [python]
def preprocess_for_test_v0(image, shape, output_height, output_width, data_format='NCHW', bbox_border=25., heatmap_sigma=1., heatmap_size=64, scope=None):
  """Preprocesses the given image for evaluation.

  Args:
    image: A `Tensor` representing an image of arbitrary size.
    output_height: The height of the image after preprocessing.
    output_width: The width of the image after preprocessing.

  Returns:
    A preprocessed image.
  """
  with tf.name_scope(scope, 'vgg_test_image', [image, output_height, output_width]):
    # Crop the central region of the image with an area containing 87.5% of
    # the original image.

    image = tf.expand_dims(image, 0)
    image = tf.image.resize_bilinear(image, [output_height, output_width], align_corners=False)
    image = tf.squeeze(image, [0])
    image.set_shape([output_height, output_width, 3])
    image = tf.cast(image, tf.float32)

    normarlized_image = _mean_image_subtraction(image, [_R_MEAN, _G_MEAN, _B_MEAN])
    if data_format == 'NCHW':
      normarlized_image = tf.transpose(normarlized_image, perm=(2, 0, 1))
    return normarlized_image/255.
//  [python]
def preprocess_for_test_raw_output(image, output_height, output_width, data_format='NCHW', scope=None):
  """Preprocesses the given image for evaluation.

  Args:
    image: A `Tensor` representing an image of arbitrary size.
    output_height: The height of the image after preprocessing.
    output_width: The width of the image after preprocessing.

  Returns:
    A preprocessed image.
  """
  with tf.name_scope(scope, 'vgg_test_image_raw_output', [image, output_height, output_width]):
    # Crop the central region of the image with an area containing 87.5% of
    # the original image.
    image = tf.image.resize_bilinear(image, [output_height, output_width], align_corners=False)
    image = tf.squeeze(image, [0])
    image.set_shape([output_height, output_width, 3])

    if config.DEBUG:
      save_image_op = tf.py_func(_save_image,
                                  [image],
                                  tf.int64, stateful=True)
      image = tf.Print(image, [save_image_op])

    image = tf.to_float(image)
    normarlized_image = _mean_image_subtraction(image, [_R_MEAN, _G_MEAN, _B_MEAN])
    if data_format == 'NCHW':
      normarlized_image = tf.transpose(normarlized_image, perm=(2, 0, 1))
    return tf.expand_dims(normarlized_image/255., 0)
//  [python]
def preprocess_for_test_raw_output(image, output_height, output_width, data_format='NCHW', scope=None):
  """Preprocesses the given image for evaluation.

  Args:
    image: A `Tensor` representing an image of arbitrary size.
    output_height: The height of the image after preprocessing.
    output_width: The width of the image after preprocessing.

  Returns:
    A preprocessed image.
  """
  with tf.name_scope(scope, 'vgg_test_image_raw_output', [image, output_height, output_width]):
    # Crop the central region of the image with an area containing 87.5% of
    # the original image.
    image = tf.image.resize_bilinear(image, [output_height, output_width], align_corners=False)
    image = tf.squeeze(image, [0])
    image.set_shape([output_height, output_width, 3])

    if config.DEBUG:
      save_image_op = tf.py_func(_save_image,
                                  [image],
                                  tf.int64, stateful=True)
      image = tf.Print(image, [save_image_op])

    image = tf.cast(image, tf.float32)
    normarlized_image = _mean_image_subtraction(image, [_R_MEAN, _G_MEAN, _B_MEAN])
    if data_format == 'NCHW':
      normarlized_image = tf.transpose(normarlized_image, perm=(2, 0, 1))
    return tf.expand_dims(normarlized_image/255., 0)
//  [python]
def test_load_all_det_checkpoint_vars(self):
    test_graph_detection = tf.Graph()
    with test_graph_detection.as_default():
      model = self._build_model(
          is_training=False,
          number_of_stages=2,
          second_stage_batch_size=6,
          num_classes=42)

      inputs_shape = (2, 20, 20, 3)
      inputs = tf.to_float(
          tf.random_uniform(inputs_shape, minval=0, maxval=255, dtype=tf.int32))
      preprocessed_inputs, true_image_shapes = model.preprocess(inputs)
      prediction_dict = model.predict(preprocessed_inputs, true_image_shapes)
      model.postprocess(prediction_dict, true_image_shapes)
      another_variable = tf.Variable([17.0], name='another_variable')  # pylint: disable=unused-variable
      var_map = model.restore_map(
          fine_tune_checkpoint_type='detection',
          load_all_detection_checkpoint_vars=True)
      self.assertIsInstance(var_map, dict)
      self.assertIn('another_variable', var_map)
//  [python]
def test_load_all_det_checkpoint_vars(self):
    test_graph_detection = tf.Graph()
    with test_graph_detection.as_default():
      model = self._build_model(
          is_training=False,
          number_of_stages=2,
          second_stage_batch_size=6,
          num_classes=42)

      inputs_shape = (2, 20, 20, 3)
      inputs = tf.cast(
          tf.random_uniform(inputs_shape, minval=0, maxval=255, dtype=tf.int32), dtype=tf.float32)
      preprocessed_inputs, true_image_shapes = model.preprocess(inputs)
      prediction_dict = model.predict(preprocessed_inputs, true_image_shapes)
      model.postprocess(prediction_dict, true_image_shapes)
      another_variable = tf.Variable([17.0], name='another_variable')  # pylint: disable=unused-variable
      var_map = model.restore_map(
          fine_tune_checkpoint_type='detection',
          load_all_detection_checkpoint_vars=True)
      self.assertIsInstance(var_map, dict)
      self.assertIn('another_variable', var_map)
//  [python]
def _apply_eig_nomean(self, eig, c, max_abs_eig):
    """tf.cond true_fn for eig_nomean."""
    eig_result = eig[:, :, c] / max_abs_eig[c]
    a = tf.pow(tf.abs(eig_result), self.pow_nomean[c])
    b = (tf.to_float(eig_result > 0) - 0.5) * 2
    eig_result = tf.multiply(a, b)
    eig_result = eig_result + self.add_nomean[c]
    eig_result = eig_result * self.mult_nomean[c]
    return eig_result
//  [python]
def _apply_eig_nomean(self, eig, c, max_abs_eig):
    """tf.cond true_fn for eig_nomean."""
    eig_result = eig[:, :, c] / max_abs_eig[c]
    a = tf.pow(tf.abs(eig_result), self.pow_nomean[c])
    b = (tf.cast(eig_result > 0, tf.float32) - 0.5) * 2
    eig_result = tf.multiply(a, b)
    eig_result = eig_result + self.add_nomean[c]
    eig_result = eig_result * self.mult_nomean[c]
    return eig_result
//  [python]
def _apply_eig_withmean(self, eig, max_abs_eig):
    """tf.cond true_fn."""
    a = tf.pow(tf.abs(eig[:, :, 0]), self.pow_withmean[0])
    b = (tf.to_float(eig[:, :, 0] > 0) - 0.5) * 2
    eig0 = tf.multiply(a, b)
    eig0 += self.add_withmean[0]
    eig0 *= self.mult_withmean[0]
    return eig0
//  [python]
def _apply_eig_withmean(self, eig, max_abs_eig):
    """tf.cond true_fn."""
    a = tf.pow(tf.abs(eig[:, :, 0]), self.pow_withmean[0])
    b = (tf.cast(eig[:, :, 0] > 0, tf.float32) - 0.5) * 2
    eig0 = tf.multiply(a, b)
    eig0 += self.add_withmean[0]
    eig0 *= self.mult_withmean[0]
    return eig0
//  [python]
def preprocess_for_eval(image, classid, shape, output_height, output_width, key_x, key_y, key_v, norm_table, data_format, category, bbox_border, heatmap_sigma, heatmap_size, resize_side, scope=None):
  """Preprocesses the given image for evaluation.

  Args:
    image: A `Tensor` representing an image of arbitrary size.
    output_height: The height of the image after preprocessing.
    output_width: The width of the image after preprocessing.
    resize_side: The smallest side of the image for aspect-preserving resizing.

  Returns:
    A preprocessed image.
  """
  with tf.name_scope(scope, 'vgg_eval_image', [image, output_height, output_width]):
    # Crop the central region of the image with an area containing 87.5% of
    # the original image.
    fkey_x, fkey_y = tf.cast(key_x, tf.float32)/tf.cast(shape[1], tf.float32), tf.cast(key_y, tf.float32)/tf.cast(shape[0], tf.float32)
    image = tf.expand_dims(image, 0)
    image = tf.image.resize_bilinear(image, [output_height, output_width], align_corners=False)
    image = tf.squeeze(image, [0])
    image.set_shape([output_height, output_width, 3])
    image = tf.to_float(image)

    ikey_x = tf.cast(tf.round(fkey_x * heatmap_size), tf.int64)
    ikey_y = tf.cast(tf.round(fkey_y * heatmap_size), tf.int64)

    targets, isvalid = draw_labelmap(ikey_x, ikey_y, heatmap_sigma, heatmap_size)

    norm_gather_ind = tf.stack([norm_table[0].lookup(classid), norm_table[1].lookup(classid)], axis=-1)

    key_x = tf.cast(tf.round(fkey_x * output_width), tf.int64)
    key_y = tf.cast(tf.round(fkey_y * output_height), tf.int64)

    norm_x, norm_y = tf.cast(tf.gather(key_x, norm_gather_ind), tf.float32), tf.cast(tf.gather(key_y, norm_gather_ind), tf.float32)
    norm_x, norm_y = tf.squeeze(norm_x), tf.squeeze(norm_y)
    norm_value = tf.pow(tf.pow(norm_x[0] - norm_x[1], 2.) + tf.pow(norm_y[0] - norm_y[1], 2.), .5)

    if config.DEBUG:
      save_image_op = tf.py_func(save_image_with_heatmap,
                                  [image, targets,
                                  config.left_right_group_map[category][0],
                                  config.left_right_group_map[category][1],
                                  config.left_right_group_map[category][2],
                                  [output_height, output_width],
                                  heatmap_size],
                                  tf.int64, stateful=True)
      with tf.control_dependencies([save_image_op]):
        normarlized_image = _mean_image_subtraction(image, [_R_MEAN, _G_MEAN, _B_MEAN])
    else:
      normarlized_image = _mean_image_subtraction(image, [_R_MEAN, _G_MEAN, _B_MEAN])

    if data_format == 'NCHW':
      normarlized_image = tf.transpose(normarlized_image, perm=(2, 0, 1))
    return normarlized_image/255., targets, key_v, isvalid, norm_value