jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def update_gradients_for_elastic_workers(self, grads):
        gradient_multiplier = 1
        if self.op == hvd.Average:
            valid_grads = get_not_none_from_list(grads)
            horovod_size = tf.cast(
                hvd.size_op()
                if int(os.environ.get("HOROVOD_ELASTIC", 0))
                else hvd.size(),
                dtype=valid_grads[0].dtype,
            )
            gradient_multiplier = (
                (horovod_size / self.global_batch_count_per_step)
                if self.average_aggregated_gradients
                else (
                    horovod_size
                    * self.backward_passes_per_step
                    / self.global_batch_count_per_step
                )
            )
        elif self.op == hvd.Sum:
            gradient_multiplier = (
                1.0 / self.backward_passes_per_step
                if self.average_aggregated_gradients
                else 1
            )

        return apply_op_to_not_none_tensors(
            tf.multiply, grads, gradient_multiplier
        )
//  [python]
def getVOCSegDataset(is_train, crop_size, voc_dir, colormap2label, max_num=None):
    """
    crop_size: (h, w)
    """
    features, labels = read_voc_images(root=voc_dir, 
                        is_train=is_train,
                        max_num=max_num)
    def _filter(imgs, crop_size):
        return [img for img in imgs if (
            img.shape[0] >= crop_size[0] and
            img.shape[1] >= crop_size[1])]
    
    def _crop(features, labels):
        features_crop = []
        labels_crop = []
        for feature, label in zip(features, labels):
            feature, label = voc_rand_crop(feature, label, 
                            height=crop_size[0],
                            width=crop_size[1])
            features_crop.append(feature)
            labels_crop.append(label)
        return features_crop, labels_crop
    
    def _normalize(feature, label):
        rgb_mean = np.array([0.485, 0.456, 0.406])
        rgb_std = np.array([0.229, 0.224, 0.225])
        
        label = voc_label_indices(label, colormap2label)
        feature = t.cast(feature, t.float32)
        feature = tf.divide(feature, 255.)
        # 我试着去掉这个反而更准
        # feature = tf.divide(tf.subtract(feature, rgb_mean), rgb_std)
        return feature, label

    features = _filter(features, crop_size)
    labels = _filter(labels, crop_size)
    features, labels = _crop(features, labels)

    print('read ' + str(len(features)) + ' valid examples')
    dataset = tf.data.Dataset.from_tensor_slices((features, labels))
    dataset = dataset.map(_normalize)
    return dataset
//  [python]
def getVOCSegDataset(is_train, crop_size, voc_dir, colormap2label, max_num=None):
    """
    crop_size: (h, w)
    """
    features, labels = read_voc_images(root=voc_dir, 
                        is_train=is_train,
                        max_num=max_num)
    def _filter(imgs, crop_size):
        return [img for img in imgs if (
            img.shape[0] >= crop_size[0] and
            img.shape[1] >= crop_size[1])]
    
    def _crop(features, labels):
        features_crop = []
        labels_crop = []
        for feature, label in zip(features, labels):
            feature, label = voc_rand_crop(feature, label, 
                            height=crop_size[0],
                            width=crop_size[1])
            features_crop.append(feature)
            labels_crop.append(label)
        return features_crop, labels_crop
    
    def _normalize(feature, label):
        rgb_mean = np.array([0.485, 0.456, 0.406])
        rgb_std = np.array([0.229, 0.224, 0.225])
        
        label = voc_label_indices(label, colormap2label)
        feature = tf.cast(feature, tf.float32)
        feature = tf.divide(feature, 255.)
        # 我试着去掉这个反而更准
        # feature = tf.divide(tf.subtract(feature, rgb_mean), rgb_std)
        return feature, label

    features = _filter(features, crop_size)
    labels = _filter(labels, crop_size)
    features, labels = _crop(features, labels)

    print('read ' + str(len(features)) + ' valid examples')
    dataset = tf.data.Dataset.from_tensor_slices((features, labels))
    dataset = dataset.map(_normalize)
    return dataset
//  [python]
def loss(self, y_true, y_pred):
        eps = sys.float_info.epsilon
        if self.keep_labels is not None:
            mask = tf.zeros(y_true.shape, dtype=tf.bool)
            for keep_label  in self.keep_labels:
                mask = tf.logical_or(mask, y_pred[...,1:2] == keep_label)

            mask = tf.cast(mask, y_pred.dtype)
        else:   # if none specified, just use all non-zero labels
            mask = tf.cast(y_pred[...,1:2] > 0, y_pred.dtype)
        pred_image = tf.math.log(y_pred[...,0:1]+eps)
        true_image = tf.math.log(y_true + eps)
        dif = tf.math.squared_difference(true_image, pred_image)
        if self.debug:
            fv = fs.Freeview(swap_batch_dim=True)
            fv.vol(true_image, name='true', opts=':linked=1:visible=0')
            fv.vol(pred_image, name='pred', opts=':linked=1')
            fv.vol(dif*mask, name='dif', opts=':linked=1:colormap=heat:locked=1')
            fv.show()
            
        lval = 10. * tf.math.divide_no_nan(tf.reduce_sum(dif * mask), tf.reduce_sum(mask))
        return lval
//  [python]
def _adjust_labels(labels, predictions):
  """Adjust the 'labels' tensor by squeezing it if needed."""
  labels = tf.to_int32(labels)
  if len(predictions.shape) == len(labels.shape):
    labels = tf.squeeze(labels, [-1])
  return labels, predictions
//  [python]
def _adjust_labels(labels, predictions):
  """Adjust the 'labels' tensor by squeezing it if needed."""
  labels = tf.cast(labels, tf.int32)
  if len(predictions.shape) == len(labels.shape):
    labels = tf.squeeze(labels, [-1])
  return labels, predictions
//  [python]
def call(self, input_ids):
    """Implements call() for the layer.

    Args:
      input_ids: int32 Tensor of shape [batch_size, seq_length].

    Returns:
      boolean Tensor of shape [batch_size, seq_length, seq_length]. [x, y, z]
      is True if for x'th instance in a batch, y'th token and z'th token are
      from the same sub sequence.
    """
    # Suppose
    # - the first token in the parent sequence is [CLS].
    # - every sequence starts from [CLS].
    # - every sequence only contains one [CLS].
    seq_start_token = input_ids[:, 0:1]
    seq_start_loc = t.cast(tf.equal(input_ids, seq_start_token), tf.int32)
    # Set different ids for different sub sequences.
    seq_ids = tf.expand_dims(tf.cumsum(seq_start_loc, -1), -1)
    return tf.equal(seq_ids, tf.transpose(seq_ids, [0, 2, 1]))
//  [python]
def call(self, input_ids):
    """Implements call() for the layer.

    Args:
      input_ids: int32 Tensor of shape [batch_size, seq_length].

    Returns:
      boolean Tensor of shape [batch_size, seq_length, seq_length]. [x, y, z]
      is True if for x'th instance in a batch, y'th token and z'th token are
      from the same sub sequence.
    """
    # Suppose
    # - the first token in the parent sequence is [CLS].
    # - every sequence starts from [CLS].
    # - every sequence only contains one [CLS].
    seq_start_token = input_ids[:, 0:1]
    seq_start_loc = tf.cast(tf.equal(input_ids, seq_start_token), tf.int32)
    # Set different ids for different sub sequences.
    seq_ids = tf.expand_dims(tf.cumsum(seq_start_loc, -1), -1)
    return tf.equal(seq_ids, tf.transpose(seq_ids, [0, 2, 1]))
//  [python]
def call(self, inputs):
    """Calls the layer with the given inputs."""
    scale = tf.cast(self._scale, inputs.dtype)
    return scale * inputs
//  [python]
def loss(labels, predictions, weights=None, from_logits=False):
  """Calculate a per-batch sparse categorical crossentropy loss.

  This loss function assumes that the predictions are post-softmax.
  Args:
    labels: The labels to evaluate against. Should be a set of integer indices
      ranging from 0 to (vocab_size-1).
    predictions: The network predictions. Should have softmax already applied.
    weights: An optional weight array of the same shape as the 'labels' array.
      If None, all examples will be used.
    from_logits: Whether the input predictions are logits.

  Returns:
    A loss scalar.

  Raises:
    RuntimeError if the passed tensors do not have the same rank.
  """
  # When using these functions with the Keras core API, we will need to squeeze
  # the labels tensor - Keras adds a spurious inner dimension.
  labels, predictions = _adjust_labels(labels, predictions)
  _validate_rank(labels, predictions, weights)

  example_losses = tf_keras.losses.sparse_categorical_crossentropy(
      labels, predictions, from_logits=from_logits)

  if weights is None:
    return tf.reduce_mean(example_losses)
  weights = tf.to_float(weights)
  return tf.math.divide_no_nan(
      tf.reduce_sum(example_losses * weights), tf.reduce_sum(weights))
//  [python]
def loss(labels, predictions, weights=None, from_logits=False):
  """Calculate a per-batch sparse categorical crossentropy loss.

  This loss function assumes that the predictions are post-softmax.
  Args:
    labels: The labels to evaluate against. Should be a set of integer indices
      ranging from 0 to (vocab_size-1).
    predictions: The network predictions. Should have softmax already applied.
    weights: An optional weight array of the same shape as the 'labels' array.
      If None, all examples will be used.
    from_logits: Whether the input predictions are logits.

  Returns:
    A loss scalar.

  Raises:
    RuntimeError if the passed tensors do not have the same rank.
  """
  # When using these functions with the Keras core API, we will need to squeeze
  # the labels tensor - Keras adds a spurious inner dimension.
  labels, predictions = _adjust_labels(labels, predictions)
  _validate_rank(labels, predictions, weights)

  example_losses = tf_keras.losses.sparse_categorical_crossentropy(
      labels, predictions, from_logits=from_logits)

  if weights is None:
    return tf.reduce_mean(example_losses)
  weights = tf.cast(weights, predictions.dtype)
  return tf.math.divide_no_nan(
      tf.reduce_sum(example_losses * weights), tf.reduce_sum(weights))
//  [python]
def scatter_labels(labels, batch_size, sequence_length):
  """Create one hot labels."""
  row_ids = labels.value_rowids()
  indices = tf.concat(
      (tf.stack((row_ids, tf.to_int32(labels.flat_values[:, 0]), tf.zeros_like(row_ids)), -1),
       tf.stack((row_ids, tf.to_int32(labels.flat_values[:, 1]), tf.ones_like(row_ids)), -1)), 0)
  one_hot_labels = tf.scatter_nd(indices,
                                 tf.ones(tf.shape(indices)[0], tf.float32),
                                 (batch_size, sequence_length, 2))
  return tf.minimum(one_hot_labels, 1.)
//  [python]
def scatter_labels(labels, batch_size, sequence_length):
  """Create one hot labels."""
  row_ids = labels.value_rowids()
  indices = tf.concat(
      (tf.stack((row_ids, tf.cast(labels.flat_values[:, 0],
                                  tf.int32), tf.zeros_like(row_ids)), -1),
       tf.stack((row_ids, tf.cast(labels.flat_values[:, 1],
                                  tf.int32), tf.ones_like(row_ids)), -1)), 0)
  one_hot_labels = tf.scatter_nd(indices,
                                 tf.ones(tf.shape(indices)[0], tf.float32),
                                 (batch_size, sequence_length, 2))
  return tf.minimum(one_hot_labels, 1.)
//  [python]
def labels_map_fn(token_ids, labels, sequence_length):
  batch_size = tf.get_static_value(labels.shape[0])
  row_lengths = labels.row_lengths()
  empty_token_index = token_ids.row_lengths() - 1
  one_hot_labels = scatter_labels(labels, batch_size, sequence_length)
  one_hot_labels += (t.cast(row_lengths == 0, t.float32)[:, tf.newaxis] *
                     tf.one_hot(empty_token_index, sequence_length))[:, :,
                                                                     tf.newaxis]
  return one_hot_labels
//  [python]
def labels_map_fn(token_ids, labels, sequence_length):
  batch_size = tf.get_static_value(labels.shape[0])
  row_lengths = labels.row_lengths()
  empty_token_index = token_ids.row_lengths() - 1
  one_hot_labels = scatter_labels(labels, batch_size, sequence_length)
  one_hot_labels += (tf.cast(row_lengths == 0, tf.float32)[:, tf.newaxis] *
                     tf.one_hot(empty_token_index, sequence_length))[:, :,
                                                                     tf.newaxis]
  return one_hot_labels
//  [python]
def load_baseline_image(filename, image_shape=None):
  """Loads baseline image and makes sure it is of the right shape.

  Args:
    filename: file name of the image to load.
    image_shape: expected shape of the image.

  Returns:
    tf.Tensor with baseline image
  """
  image_path = tf.compat.v1.resource_loader.get_path_to_datafile(
      os.path.join("test_data", filename))
  file = tf.io.read_file(image_path)
  baseline_image = tf.to_float(tf.image.decode_image(file)) / 255.0
  baseline_image = tf.expand_dims(baseline_image, axis=0)
  if image_shape is not None:
    # Graph-mode requires image shape to be known in advance.
    baseline_image = tf.ensure_shape(baseline_image, image_shape)
  return baseline_image
//  [python]
def load_baseline_image(filename, image_shape=None):
  """Loads baseline image and makes sure it is of the right shape.

  Args:
    filename: file name of the image to load.
    image_shape: expected shape of the image.

  Returns:
    tf.Tensor with baseline image
  """
  image_path = tf.compat.v1.resource_loader.get_path_to_datafile(
      os.path.join("test_data", filename))
  file = tf.io.read_file(image_path)
  baseline_image = tf.cast(tf.image.decode_image(file), tf.float32) / 255.0
  baseline_image = tf.expand_dims(baseline_image, axis=0)
  if image_shape is not None:
    # Graph-mode requires image shape to be known in advance.
    baseline_image = tf.ensure_shape(baseline_image, image_shape)
  return baseline_image
//  [python]
def call(self, inputs):
    """Calls the layer with the given inputs."""
    assert inputs.shape.rank in (3, 4, 5)
    frames = tf.shape(inputs)[1]
    pre_logits = inputs / tf.sqrt(tf.to_float(frames))
    activations = tf.nn.softmax(pre_logits, axis=1)
    outputs = inputs * activations
    return outputs
//  [python]
def call(self, inputs):
    """Calls the layer with the given inputs."""
    assert inputs.shape.rank in (3, 4, 5)
    frames = tf.shape(inputs)[1]
    pre_logits = inputs / tf.sqrt(tf.cast(frames, inputs.dtype))
    activations = tf.nn.softmax(pre_logits, axis=1)
    outputs = inputs * activations
    return outputs
//  [python]
def call(
      self,
      y_true: tf.Tensor,
      y_pred: tf.Tensor):
    """Compute softmax loss with batch labels as negatives.

    Args:
      y_true: The weights for each example in the batch with shape [batch_size]
        or [batch_size, 1].
      y_pred: When expect_embedding is False, `y_pred` is the pre-calculated
        batch similarities matrix with shape [batch_size, batch_size]; when
        `expected_embedding` is True, `y_pred` is the output batch context and
        label embedddings of dual encoder model [context_embedding,
        label_embedding].

    Returns:
      The softmax loss with batch labels as negatives. Optionally add
      spreadout regularizations with given scaling lambdas.
    """

    context_embedding, label_embedding, similarities = (
        utils.get_embeddings_and_similarities(
            y_pred, y_true, self.expect_embeddings, self.normalization_fn))

    # Use the diagonal elements of similarities as labels for each row.
    curr_batch_size = tf.shape(similarities)[0]
    label_indices = self.label_indices_fn(curr_batch_size)
    losses = tf.nn.sparse_softmax_cross_entropy_with_logits(
        labels=label_indices, logits=similarities)

    # Perform a weighted sum, then divide by the sum of weights to get the
    # weighted average of the per-example losses. Flatten weights first in case
    # they have shape (batch_size, 1).
    y_true = tf.reshape(y_true, [-1])
    loss = tf.math.divide_no_nan(
        tf.reduce_sum(losses * y_true), tf.reduce_sum(y_true))

    # Apply spreadout if needed.
    loss = _update_loss_with_spreadout_loss(
        loss=loss,
        context_embedding=context_embedding,
        label_embedding=label_embedding,
        similarities=similarities,
        spreadout_context_lambda=self.spreadout_context_lambda,
        spreadout_label_lambda=self.spreadout_label_lambda,
        spreadout_cross_lambda=self.spreadout_cross_lambda,
        label_indices=label_indices)

    return loss
//  [python]
def call(
      self,
      y_true: tf.Tensor,
      y_pred: tf.Tensor):
    """Compute softmax loss with batch labels as negatives.

    Args:
      y_true: The weights for each example in the batch with shape [batch_size]
        or [batch_size, 1].
      y_pred: When expect_embedding is False, `y_pred` is the pre-calculated
        batch similarities matrix with shape [batch_size, batch_size]; when
        `expected_embedding` is True, `y_pred` is the output batch context and
        label embedddings of dual encoder model [context_embedding,
        label_embedding].

    Returns:
      The softmax loss with batch labels as negatives. Optionally add
      spreadout regularizations with given scaling lambdas.
    """

    context_embedding, label_embedding, similarities = (
        utils.get_embeddings_and_similarities(
            y_pred, y_true, self.expect_embeddings, self.normalization_fn))

    # Use the diagonal elements of similarities as labels for each row.
    curr_batch_size = tf.shape(similarities)[0]
    label_indices = self.label_indices_fn(curr_batch_size)
    losses = tf.nn.sparse_softmax_cross_entropy_with_logits(
        labels=label_indices, logits=similarities)

    # Perform a weighted sum, then divide by the sum of weights to get the
    # weighted average of the per-example losses. Flatten weights first in case
    # they have shape (batch_size, 1).
    y_true = tf.cast(tf.reshape(y_true, [-1]), losses.dtype)
    loss = tf.math.divide_no_nan(
        tf.reduce_sum(losses * y_true), tf.reduce_sum(y_true))

    # Apply spreadout if needed.
    loss = _update_loss_with_spreadout_loss(
        loss=loss,
        context_embedding=context_embedding,
        label_embedding=label_embedding,
        similarities=similarities,
        spreadout_context_lambda=self.spreadout_context_lambda,
        spreadout_label_lambda=self.spreadout_label_lambda,
        spreadout_cross_lambda=self.spreadout_cross_lambda,
        label_indices=label_indices)

    return loss
//  [python]
def pyramid_feature_fusion(inputs, target_level):
  """Fuses all feature maps in the feature pyramid at the target level.

  Args:
    inputs: A dictionary containing the feature pyramid. The size of the input
      tensor needs to be fixed.
    target_level: An `int` of the target feature level for feature fusion.

  Returns:
    A `float` `tf.Tensor` of shape [batch_size, feature_height, feature_width,
      feature_channel].
  """
  # Convert keys to int.
  pyramid_feats = {int(k): v for k, v in inputs.items()}
  min_level = min(pyramid_feats.keys())
  max_level = max(pyramid_feats.keys())
  resampled_feats = []

  for l in range(min_level, max_level + 1):
    if l == target_level:
      resampled_feats.append(pyramid_feats[l])
    else:
      feat = pyramid_feats[l]
      target_size = list(feat.shape[1:3])
      target_size[0] *= 2**(l - target_level)
      target_size[1] *= 2**(l - target_level)
      # Casts feat to float32 so the resize op can be run on TPU.
      feat = tf.to_float(feat)
      feat = tf.image.resize(
          feat, size=target_size, method=tf.image.ResizeMethod.BILINEAR)
      # Casts it back to be compatible with the rest opetations.
      feat = tf.to_int32(feat)
      resampled_feats.append(feat)

  return tf.math.add_n(resampled_feats)
//  [python]
def pyramid_feature_fusion(inputs, target_level):
  """Fuses all feature maps in the feature pyramid at the target level.

  Args:
    inputs: A dictionary containing the feature pyramid. The size of the input
      tensor needs to be fixed.
    target_level: An `int` of the target feature level for feature fusion.

  Returns:
    A `float` `tf.Tensor` of shape [batch_size, feature_height, feature_width,
      feature_channel].
  """
  # Convert keys to int.
  pyramid_feats = {int(k): v for k, v in inputs.items()}
  min_level = min(pyramid_feats.keys())
  max_level = max(pyramid_feats.keys())
  resampled_feats = []

  for l in range(min_level, max_level + 1):
    if l == target_level:
      resampled_feats.append(pyramid_feats[l])
    else:
      feat = pyramid_feats[l]
      target_size = list(feat.shape[1:3])
      target_size[0] *= 2**(l - target_level)
      target_size[1] *= 2**(l - target_level)
      # Casts feat to float32 so the resize op can be run on TPU.
      feat = tf.cast(feat, tf.float32)
      feat = tf.image.resize(
          feat, size=target_size, method=tf.image.ResizeMethod.BILINEAR)
      # Casts it back to be compatible with the rest opetations.
      feat = tf.cast(feat, pyramid_feats[l].dtype)
      resampled_feats.append(feat)

  return tf.math.add_n(resampled_feats)
//  [python]
def compute_wer(decode, target, dtype=tf.float32):
        decode = tf.strings.split(decode)
        target = tf.strings.split(target)
        distances = tf.to_float(tf.edit_distance(decode.to_sparse(), target.to_sparse(), normalize=False))  # [B]
        lengths = tf.to_float(target.row_lengths(axis=1))  # [B]
        return distances, lengths
//  [python]
def compute_wer(decode, target, dtype=tf.float32):
        decode = tf.strings.split(decode)
        target = tf.strings.split(target)
        distances = tf.cast(tf.edit_distance(decode.to_sparse(), target.to_sparse(), normalize=False), dtype)  # [B]
        lengths = tf.cast(target.row_lengths(axis=1), dtype)  # [B]
        return distances, lengths
//  [python]
def compute_cer(decode, target, dtype=tf.float32):
        decode = tf.strings.bytes_split(decode)  # [B, N]
        target = tf.strings.bytes_split(target)  # [B, M]
        distances = tf.to_float(tf.edit_distance(decode.to_sparse(), target.to_sparse(), normalize=False)) # [B]
        lengths = tf.to_float(target.row_lengths(axis=1))  # [B]
        return distances, lengths
//  [python]
def compute_cer(decode, target, dtype=tf.float32):
        decode = tf.strings.bytes_split(decode)  # [B, N]
        target = tf.strings.bytes_split(target)  # [B, M]
        distances = tf.cast(tf.edit_distance(decode.to_sparse(), target.to_sparse(), normalize=False), dtype)  # [B]
        lengths = tf.cast(target.row_lengths(axis=1), dtype)  # [B]
        return distances, lengths
//  [python]
def call(self,
           inputs: tf.Tensor,
           training: Optional[bool] = None) -> tf.Tensor:
    if training is None:
      training = tf_keras.backend.learning_phase()
    result = []
    for i, layers in enumerate(self.aspp_layers):
      x = inputs
      for layer in layers:
        # Apply layers sequentially.
        x = layer(x, training=training)
      x = self._activation_fn(x)

      # Apply resize layer to the end of the last set of layers.
      if i == len(self.aspp_layers) - 1:
        x = self._resizing_layer(x)

      result.append(tf.cast(x, inputs.dtype))
    x = self._concat_layer(result)
    for layer in self._projection:
      x = layer(x, training=training)
    x = self._activation_fn(x)
    return self._dropout_layer(x)
//  [python]
def resize(x, scale=2, to_shape=None, align_corners=True, dynamic=False,
           func=tf.image.resize_bilinear, name='resize'):
    if dynamic:
        xs = t.cast(t.shape(x), t.float32)
        new_xs = [t.cast(xs[1]*scale, t.int32),
                  t.cast(xs[2]*scale, t.int32)]
    else:
        xs = x.get_shape().as_list()
        new_xs = [int(xs[1]*scale), int(xs[2]*scale)]
    with tf.variable_scope(name):
        if to_shape is None:
            x = func(x, new_xs, align_corners=align_corners)
        else:
            x = func(x, [to_shape[0], to_shape[1]],
                     align_corners=align_corners)
    return x
//  [python]
def resize(x, scale=2, to_shape=None, align_corners=True, dynamic=False,
           func=tf.image.resize_bilinear, name='resize'):
    if dynamic:
        xs = tf.cast(tf.shape(x), tf.float32)
        new_xs = [tf.cast(xs[1]*scale, tf.int32),
                  tf.cast(xs[2]*scale, tf.int32)]
    else:
        xs = x.get_shape().as_list()
        new_xs = [int(xs[1]*scale), int(xs[2]*scale)]
    with tf.variable_scope(name):
        if to_shape is None:
            x = func(x, new_xs, align_corners=align_corners)
        else:
            x = func(x, [to_shape[0], to_shape[1]],
                     align_corners=align_corners)
    return x
//  [python]
def _parse_data(record):
    feature_description = {
        "image": tf.io.FixedLenFeature([28, 28], tf.float32),
        "label": tf.io.FixedLenFeature([1], tf.int64),
    }
    r = tf.io.parse_single_example(record, feature_description)
    features = {
        "image": tf.math.divide(tf.cast(r["image"], tf.float32), 255.0)
    }
    return features, tf.to_int32(r["label"])
//  [python]
def _parse_data(record):
    feature_description = {
        "image": tf.io.FixedLenFeature([28, 28], tf.float32),
        "label": tf.io.FixedLenFeature([1], tf.int64),
    }
    r = tf.io.parse_single_example(record, feature_description)
    features = {
        "image": tf.math.divide(tf.cast(r["image"], tf.float32), 255.0)
    }
    return features, tf.cast(r["label"], tf.int32)
//  [python]
def _positional_encoding(self,
                           num_positions: Union[int, tf.Tensor],
                           hidden_size: Union[int, tf.Tensor],
                           start_position: Union[int, tf.Tensor] = 0,
                           dtype: str = 'float32') -> tf.Tensor:
    """Creates a sequence of sinusoidal positional encoding vectors.

    Args:
      num_positions: the total number of positions (frames).
      hidden_size: the number of channels used for the hidden vectors.
      start_position: the start position.
      dtype: the dtype of the output tensor.

    Returns:
      The positional encoding tensor with shape [num_positions, hidden_size].
    """
    if isinstance(start_position, tf.Tensor) and start_position.shape.rank == 1:
      start_position = start_position[0]

    # Calling `tf.range` with `dtype=tf.bfloat16` results in an error,
    # so we cast afterward.
    positions = tf.range(start_position, start_position + num_positions)
    positions = tf.cast(positions, dtype)[:, tf.newaxis]
    idx = tf.range(hidden_size)[tf.newaxis, :]

    power = tf.cast(2 * (idx // 2), dtype)
    power /= tf.cast(hidden_size, dtype)
    angles = 1. / tf.math.pow(10_000., power)
    radians = positions * angles

    sin = tf.math.sin(radians[:, 0::2])
    cos = tf.math.cos(radians[:, 1::2])
    pos_encoding = tf.concat([sin, cos], axis=-1)

    return pos_encoding
//  [python]
def call(
      self,
      inputs: tf.Tensor,
      states: Optional[States] = None,
      output_states: bool = True,
  ) -> Union[tf.Tensor, Tuple[tf.Tensor, States]]:
    """Calls the layer with the given inputs.

    Args:
      inputs: An input `tf.Tensor`.
      states: A `dict` of states such that, if any of the keys match for this
        layer, will overwrite the contents of the buffer(s). Expected keys
        include `state_prefix + '_pos_enc_frame_count'`.
      output_states: A `bool`. If True, returns the output tensor and output
        states. Returns just the output tensor otherwise.

    Returns:
      An output `tf.Tensor` (and optionally the states if `output_states=True`).

    Raises:
      ValueError: If using 'channels_first' data format.
    """
    states = dict(states) if states is not None else {}

    # Keep a count of frames encountered across input iterations in
    # num_frames to be able to accurately update the positional encoding.
    num_frames = tf.shape(inputs)[1]
    frame_count = tf.cast(states.get(self._frame_count_name, [0]), tf.int32)
    states[self._frame_count_name] = frame_count + num_frames

    if self._cache_encoding:
      pos_encoding = self._pos_encoding
    else:
      pos_encoding = self._get_pos_encoding(
          tf.shape(inputs), frame_count=frame_count)
    pos_encoding = tf.cast(pos_encoding, inputs.dtype)
    pos_encoding = self._rezero(pos_encoding)
    outputs = inputs + pos_encoding

    return (outputs, states) if output_states else outputs
//  [python]
def call(self, y_true, y_pred):
    """Invokes the `FocalLoss`.

    Args:
      y_true: A tensor of size [batch, num_anchors, num_classes]
      y_pred: A tensor of size [batch, num_anchors, num_classes]

    Returns:
      Summed loss float `Tensor`.
    """
    with tf.name_scope('focal_loss'):
      y_true = tf.cast(y_true, dtype=tf.float32)
      y_pred = tf.cast(y_pred, dtype=tf.float32)
      positive_label_mask = tf.equal(y_true, 1.0)
      cross_entropy = (
          tf.nn.sigmoid_cross_entropy_with_logits(labels=y_true, logits=y_pred))
      probs = tf.sigmoid(y_pred)
      probs_gt = tf.where(positive_label_mask, probs, 1.0 - probs)
      # With small gamma, the implementation could produce NaN during back prop.
      modulator = tf.pow(1.0 - probs_gt, self._gamma)
      loss = modulator * cross_entropy
      weighted_loss = tf.where(positive_label_mask, self._alpha * loss,
                               (1.0 - self._alpha) * loss)

    return weighted_loss
//  [python]
def build_losses(self, outputs, labels, aux_losses=None):
    """Builds DETR losses."""
    targets = labels['targets']
    weights = labels['weights']

    targets = tf.one_hot(targets, self._task_config.model.vocab_size)

    loss = tf_keras.losses.CategoricalCrossentropy(
        from_logits=True, reduction=tf_keras.losses.Reduction.NONE
    )(targets, outputs)

    weights = tf.cast(weights, loss.dtype)
    loss = tf.reduce_sum(loss * weights) / tf.reduce_sum(weights)

    aux_losses = tf.add_n(aux_losses) if aux_losses else 0.0

    total_loss = loss + aux_losses
    return total_loss
//  [python]
def _padded_cross_entropy_loss(logits, labels, smoothing, vocab_size):
  """Calculate cross entropy loss while ignoring padding.

  Args:
    logits: Tensor of size [batch_size, length_logits, vocab_size]
    labels: Tensor of size [batch_size, length_labels]
    smoothing: Label smoothing constant, used to determine the on and off values
    vocab_size: int size of the vocabulary

  Returns:
    Returns the cross entropy loss and weight tensors: float32 tensors with
      shape [batch_size, max(length_logits, length_labels)]
  """
  logits, labels = _pad_tensors_to_same_length(logits, labels)

  # Calculate smoothing cross entropy
  confidence = 1.0 - smoothing
  low_confidence = (1.0 - confidence) / t.cast(vocab_size - 1, t.float32)
  soft_targets = t.one_hot(
      t.cast(labels, t.int32),
      depth=vocab_size,
      on_value=confidence,
      off_value=low_confidence)
  xentropy = t.nn.softmax_cross_entropy_with_logits(
      logits=logits, labels=soft_targets)

  # Calculate the best (lowest) possible value of cross entropy, and
  # subtract from the cross entropy loss.
  normalizing_constant = -(
      confidence * t.math.log(confidence) + t.cast(vocab_size - 1, t.float32)
      * low_confidence * t.math.log(low_confidence + 1e-20))
  xentropy -= normalizing_constant

  weights = t.cast(t.not_equal(labels, 0), t.float32)
  return xentropy * weights, weights
//  [python]
def _padded_cross_entropy_loss(logits, labels, smoothing, vocab_size):
  """Calculate cross entropy loss while ignoring padding.

  Args:
    logits: Tensor of size [batch_size, length_logits, vocab_size]
    labels: Tensor of size [batch_size, length_labels]
    smoothing: Label smoothing constant, used to determine the on and off values
    vocab_size: int size of the vocabulary

  Returns:
    Returns the cross entropy loss and weight tensors: float32 tensors with
      shape [batch_size, max(length_logits, length_labels)]
  """
  logits, labels = _pad_tensors_to_same_length(logits, labels)

  # Calculate smoothing cross entropy
  confidence = 1.0 - smoothing
  low_confidence = (1.0 - confidence) / tf.cast(vocab_size - 1, tf.float32)
  soft_targets = tf.one_hot(
      tf.cast(labels, tf.int32),
      depth=vocab_size,
      on_value=confidence,
      off_value=low_confidence)
  xentropy = tf.nn.softmax_cross_entropy_with_logits(
      logits=logits, labels=soft_targets)

  # Calculate the best (lowest) possible value of cross entropy, and
  # subtract from the cross entropy loss.
  normalizing_constant = -(
      confidence * tf.math.log(confidence) + tf.cast(vocab_size - 1, tf.float32)
      * low_confidence * tf.math.log(low_confidence + 1e-20))
  xentropy -= normalizing_constant

  weights = tf.cast(tf.not_equal(labels, 0), tf.float32)
  return xentropy * weights, weights
//  [python]
def _prepare_image_and_label(
      self, data: Dict[str, Any]) -> Tuple[tf.Tensor, tf.Tensor]:
    """Prepares normalized image and label."""
    image = tf.io.decode_raw(data[self._image_field_key],
                             tf.as_dtype(tf.float32))
    label = tf.io.decode_raw(data[self._label_field_key],
                             tf.as_dtype(self._label_dtype))
    image_size = list(self._input_size) + [self._num_channels]
    image = tf.reshape(image, image_size)
    label_size = list(self._input_size) + [self._num_classes]
    label = tf.reshape(label, label_size)

    image = tf.to_float(image)
    label = tf.to_float(label)

    # TPU doesn't support tf.int64 well, use tf.int32 directly.
    if label.dtype == tf.int64:
      label = tf.to_int32(label)
    return image, label
//  [python]
def _prepare_image_and_label(
      self, data: Dict[str, Any]) -> Tuple[tf.Tensor, tf.Tensor]:
    """Prepares normalized image and label."""
    image = tf.io.decode_raw(data[self._image_field_key],
                             tf.as_dtype(tf.float32))
    label = tf.io.decode_raw(data[self._label_field_key],
                             tf.as_dtype(self._label_dtype))
    image_size = list(self._input_size) + [self._num_channels]
    image = tf.reshape(image, image_size)
    label_size = list(self._input_size) + [self._num_classes]
    label = tf.reshape(label, label_size)

    image = tf.cast(image, dtype=self._dtype)
    label = tf.cast(label, dtype=self._dtype)

    # TPU doesn't support tf.int64 well, use tf.int32 directly.
    if label.dtype == tf.int64:
      label = tf.cast(label, dtype=tf.int32)
    return image, label
//  [python]
def _parse_train_data(self, data: Dict[str,
                                         Any]) -> Tuple[tf.Tensor, tf.Tensor]:
    """Parses data for training and evaluation."""
    image, labels = self._prepare_image_and_label(data)
    # Cast image as self._dtype
    image = tf.to_float(image)

    return image, labels
//  [python]
def _parse_train_data(self, data: Dict[str,
                                         Any]) -> Tuple[tf.Tensor, tf.Tensor]:
    """Parses data for training and evaluation."""
    image, labels = self._prepare_image_and_label(data)
    # Cast image as self._dtype
    image = tf.cast(image, dtype=self._dtype)

    return image, labels
//  [python]
def _parse_eval_data(self, data: Dict[str,
                                        Any]) -> Tuple[tf.Tensor, tf.Tensor]:
    """Parses data for training and evaluation."""
    image, labels = self._prepare_image_and_label(data)
    # Cast image as self._dtype
    image = tf.cast(image, dtype=self._dtype)

    return image, labels
//  [python]
def train_step(self, inputs, model, optimizer, metrics=None):
    """Does forward and backward.

    Args:
      inputs: a dictionary of input tensors.
      model: the model, forward pass definition.
      optimizer: the optimizer for this training step.
      metrics: a nested structure of metrics objects.

    Returns:
      A dictionary of logs.
    """
    features, labels = inputs
    num_replicas = tf.distribute.get_strategy().num_replicas_in_sync

    with tf.GradientTape() as tape:
      _, outputs = model(features, labels['inputs'], training=True)
      outputs = tf.nest.map_structure(lambda x: tf.cast(x, tf.float32), outputs)

      loss = self.build_losses(
          outputs=outputs, labels=labels, aux_losses=model.losses
      )
      scaled_loss = loss / num_replicas

      # For mixed_precision policy, when LossScaleOptimizer is used, loss is
      # scaled for numerical stability.
      if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
        scaled_loss = optimizer.get_scaled_loss(scaled_loss)

    tvars = model.trainable_variables
    grads = tape.gradient(scaled_loss, tvars)
    # Scales back gradient when LossScaleOptimizer is used.
    if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
      grads = optimizer.get_unscaled_gradients(grads)
    optimizer.apply_gradients(list(zip(grads, tvars)))

    # Trainer class handles loss metric for you.
    logs = {self.loss: loss}

    all_losses = {
        'loss': loss,
    }

    # Metric results will be added to logs for you.
    if metrics:
      for m in metrics:
        m.update_state(all_losses[m.name])
    return logs
//  [python]
def __parser__(self, example_proto):
        """ This function parses a single datum

        :param example_proto:
        :return:
        """
        # configure feature and label length
        # It is crucial that for tf.string, the length is not specified, as the data is stored as a single string!
        x_config = tf.FixedLenFeature([], tf.string) \
            if self.x_dtype == tf.string else tf.FixedLenFeature([self.num_features], self.x_dtype)
        if self.num_labels == 0:
            proto_config = {'x': x_config}
        else:
            y_config = tf.FixedLenFeature([], tf.string) \
                if self.y_dtype == tf.string else tf.FixedLenFeature([self.num_labels], self.y_dtype)
            proto_config = {'x': x_config, 'y': y_config}

        # decode examples
        datum = tf.parse_single_example(example_proto, features=proto_config)
        if self.x_dtype == tf.string:  # if input is string / bytes, decode it to float32
            # first decode data to uint8, as data is stored in this way
            datum['x'] = tf.decode_raw(datum['x'], tf.uint8)
            # then cast data to tf.float32
            datum['x'] = tf.to_float(datum['x'])
            # cannot use string_to_number as there is only one string for a whole sample
            # datum['x'] = tf.strings.to_number(datum['x'], tf.float32)  # this results in possibly a large number

        # return data
        if 'y' in datum:
            # y can be present in many ways:
            # 1. a single integer, which requires y to be int32 or int64 (e.g, used in tf.gather in cbn)
            # 2. num-class bool/integer/float variables. This form is more flexible as it allows multiple classes and
            # prior probabilities as targets
            # 3. float variables in regression problem.
            # but...
            # y is stored as int (for case 1), string (for other int cases), or float (for float cases)
            # in the case of tf.string and tf.int64, convert to to int32
            if self.y_dtype == tf.string:
                # avoid using string labels like 'cat', 'dog', use integers instead
                datum['y'] = tf.decode_raw(datum['y'], tf.uint8)
                datum['y'] = tf.to_int32(datum['y'])
            if self.y_dtype == tf.int64:
                datum['y'] = tf.to_int32(datum['y'])
            return datum['x'], datum['y']
        else:
            return datum['x']
//  [python]
def __parser__(self, example_proto):
        """ This function parses a single datum

        :param example_proto:
        :return:
        """
        # configure feature and label length
        # It is crucial that for tf.string, the length is not specified, as the data is stored as a single string!
        x_config = tf.FixedLenFeature([], tf.string) \
            if self.x_dtype == tf.string else tf.FixedLenFeature([self.num_features], self.x_dtype)
        if self.num_labels == 0:
            proto_config = {'x': x_config}
        else:
            y_config = tf.FixedLenFeature([], tf.string) \
                if self.y_dtype == tf.string else tf.FixedLenFeature([self.num_labels], self.y_dtype)
            proto_config = {'x': x_config, 'y': y_config}

        # decode examples
        datum = tf.parse_single_example(example_proto, features=proto_config)
        if self.x_dtype == tf.string:  # if input is string / bytes, decode it to float32
            # first decode data to uint8, as data is stored in this way
            datum['x'] = tf.decode_raw(datum['x'], tf.uint8)
            # then cast data to tf.float32
            datum['x'] = tf.cast(datum['x'], tf.float32)
            # cannot use string_to_number as there is only one string for a whole sample
            # datum['x'] = tf.strings.to_number(datum['x'], tf.float32)  # this results in possibly a large number

        # return data
        if 'y' in datum:
            # y can be present in many ways:
            # 1. a single integer, which requires y to be int32 or int64 (e.g, used in tf.gather in cbn)
            # 2. num-class bool/integer/float variables. This form is more flexible as it allows multiple classes and
            # prior probabilities as targets
            # 3. float variables in regression problem.
            # but...
            # y is stored as int (for case 1), string (for other int cases), or float (for float cases)
            # in the case of tf.string and tf.int64, convert to to int32
            if self.y_dtype == tf.string:
                # avoid using string labels like 'cat', 'dog', use integers instead
                datum['y'] = tf.decode_raw(datum['y'], tf.uint8)
                datum['y'] = tf.cast(datum['y'], tf.int32)
            if self.y_dtype == tf.int64:
                datum['y'] = tf.cast(datum['y'], tf.int32)
            return datum['x'], datum['y']
        else:
            return datum['x']
//  [python]
def main(_):

    # Get image's height and width.
    height = 0
    width = 0
    with open(FLAGS.image_file, 'rb') as img:
        with tf.Session().as_default() as sess:
            if FLAGS.image_file.lower().endswith('png'):
                image = sess.run(tf.image.decode_png(img.read()))
            else:
                image = sess.run(tf.image.decode_jpeg(img.read()))
            height = image.shape[0]
            width = image.shape[1]
    tf.logging.info('Image size: %dx%d' % (width, height))

    with tf.Graph().as_default():
        with tf.Session().as_default() as sess:

            # Read image data.
            image_preprocessing_fn, _ = preprocessing_factory.get_preprocessing(
                FLAGS.loss_model,
                is_training=False)
            image = reader.get_image(FLAGS.image_file, height, width, image_preprocessing_fn)

            # Add batch dimension
            image = tf.expand_dims(image, 0)

            generated = model.net(image, training=False)
            generated = tf.to_int64(generated)

            # Remove batch dimension
            generated = tf.squeeze(generated, [0])

            # Restore model variables.
            saver = tf.train.Saver(tf.global_variables(), write_version=tf.train.SaverDef.V1)
            sess.run([tf.global_variables_initializer(), tf.local_variables_initializer()])
            # Use absolute path
            FLAGS.model_file = os.path.abspath(FLAGS.model_file)
            saver.restore(sess, FLAGS.model_file)

            # Make sure 'generated' directory exists.
            generated_file = 'generated/res.jpg'
            if os.path.exists('generated') is False:
                os.makedirs('generated')

            # Generate and write image data to file.
            with open(generated_file, 'wb') as img:
                start_time = time.time()
                img.write(sess.run(tf.image.encode_jpeg(generated)))
                end_time = time.time()
                tf.logging.info('Elapsed time: %fs' % (end_time - start_time))

                tf.logging.info('Done. Please check %s.' % generated_file)
//  [python]
def main(_):

    # Get image's height and width.
    height = 0
    width = 0
    with open(FLAGS.image_file, 'rb') as img:
        with tf.Session().as_default() as sess:
            if FLAGS.image_file.lower().endswith('png'):
                image = sess.run(tf.image.decode_png(img.read()))
            else:
                image = sess.run(tf.image.decode_jpeg(img.read()))
            height = image.shape[0]
            width = image.shape[1]
    tf.logging.info('Image size: %dx%d' % (width, height))

    with tf.Graph().as_default():
        with tf.Session().as_default() as sess:

            # Read image data.
            image_preprocessing_fn, _ = preprocessing_factory.get_preprocessing(
                FLAGS.loss_model,
                is_training=False)
            image = reader.get_image(FLAGS.image_file, height, width, image_preprocessing_fn)

            # Add batch dimension
            image = tf.expand_dims(image, 0)

            generated = model.net(image, training=False)
            generated = tf.cast(generated, tf.uint8)

            # Remove batch dimension
            generated = tf.squeeze(generated, [0])

            # Restore model variables.
            saver = tf.train.Saver(tf.global_variables(), write_version=tf.train.SaverDef.V1)
            sess.run([tf.global_variables_initializer(), tf.local_variables_initializer()])
            # Use absolute path
            FLAGS.model_file = os.path.abspath(FLAGS.model_file)
            saver.restore(sess, FLAGS.model_file)

            # Make sure 'generated' directory exists.
            generated_file = 'generated/res.jpg'
            if os.path.exists('generated') is False:
                os.makedirs('generated')

            # Generate and write image data to file.
            with open(generated_file, 'wb') as img:
                start_time = time.time()
                img.write(sess.run(tf.image.encode_jpeg(generated)))
                end_time = time.time()
                tf.logging.info('Elapsed time: %fs' % (end_time - start_time))

                tf.logging.info('Done. Please check %s.' % generated_file)
//  [python]
def pixel_nms(heat,kernel=[3,3],epsilon=1e-8):
        hmax=tf.nn.max_pool(heat,ksize=[1]+kernel+[1],strides=[1,1,1,1],padding='SAME')
        mask=tf.less_equal(tf.abs(hmax-heat),epsilon)
        mask = t.cast(mask,tf.float32)
        return mask*heat
//  [python]
def pixel_nms(heat,kernel=[3,3],epsilon=1e-8):
        hmax=tf.nn.max_pool(heat,ksize=[1]+kernel+[1],strides=[1,1,1,1],padding='SAME')
        mask=tf.less_equal(tf.abs(hmax-heat),epsilon)
        mask = tf.cast(mask,tf.float32)
        return mask*heat
//  [python]
def _topk(scores,K=100):
        B,H,W,C = wmlt.combined_static_and_dynamic_shape(scores)
        scores = tf.reshape(scores,[B,-1])
        topk_scores,topk_inds = tf.nn.top_k(scores,k=K)
        topk_classes = topk_inds%C
        topk_inds = topk_inds//C
        topk_ys = tf.to_float(topk_inds//W)
        topk_xs = tf.to_float(topk_inds%W)
        return topk_scores,topk_inds,topk_classes,topk_ys,topk_xs
//  [python]
def _topk(scores,K=100):
        B,H,W,C = wmlt.combined_static_and_dynamic_shape(scores)
        scores = tf.reshape(scores,[B,-1])
        topk_scores,topk_inds = tf.nn.top_k(scores,k=K)
        topk_classes = topk_inds%C
        topk_inds = topk_inds//C
        topk_ys = tf.cast(topk_inds//W,tf.float32)
        topk_xs = tf.cast(topk_inds%W,tf.float32)
        return topk_scores,topk_inds,topk_classes,topk_ys,topk_xs
//  [python]
def preprocess(self,image):
        with tf.name_scope('image_preprocess'):
            if image.dtype.base_dtype != tf.float32:
                image = tf.to_float(image)

            image=image/255.
        return image
//  [python]
def preprocess(self,image):
        with tf.name_scope('image_preprocess'):
            if image.dtype.base_dtype != tf.float32:
                image = tf.cast(image, tf.float32)

            image=image/255.
        return image