jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def res_net(n_examples,
            input_shape,
            num_classes,
            batchnorm=False,
            variational='full'):
  """Wrapper for build_resnet_v1.

  Args:
    n_examples (int): number of training points.
    input_shape (list): input shape.
    num_classes (int): number of classes (CIFAR10 has 10).
    batchnorm (bool): use of batchnorm layers.
    variational (str): 'none', 'hybrid', 'full'. whether to use variational
      inference for zero, some, or all layers.

  Returns:
      model (Model): Keras model instance whose output is a
        tfp.distributions.Categorical distribution.
  """
  inputs = tf.keras.layers.Input(shape=input_shape)
  x = build_resnet_v1(
      inputs,
      depth=20,
      variational=variational,
      batchnorm=batchnorm,
      n_examples=n_examples)

  p_fn, q_fn = mean_field_fn(empirical_bayes=True)

  def normalized_kl_fn(q, p, _):
    return tfp.distributions.kl_divergence(q, p) / tf.cast(n_examples, tf.float32)

  logits = tfp.layers.DenseLocalReparameterization(
      num_classes,
      kernel_prior_fn=p_fn,
      kernel_posterior_fn=q_fn,
      bias_prior_fn=p_fn,
      bias_posterior_fn=q_fn,
      kernel_divergence_fn=normalized_kl_fn,
      bias_divergence_fn=normalized_kl_fn)(x)
  outputs = tf.keras.layers.Lambda(lambda x: ed.Categorical(logits=x))(logits)
  return tf.keras.models.Model(inputs=inputs, outputs=outputs)
//  [python]
def select_fg_for_masks(class_targets, box_targets, boxes,
                        proposal_to_label_map, max_num_fg=128):
  """Selects the fore ground objects for mask branch during training.

  Args:
    class_targets: a tensor of shape [batch_size, num_boxes]  representing the
      class label for each box.
    box_targets: a tensor with a shape of [batch_size, num_boxes, 4]. The tensor
      contains the ground truth pixel coordinates of the scaled images for each
      roi.
    boxes: A 3-D Tensor of shape [batch_size, num_boxes, 4]. Each row
      represents a box with [y1, x1, y2, x2] in un-normalized coordinates.
    proposal_to_label_map: a tensor with a shape of [batch_size, num_boxes].
      This tensor keeps the mapping between proposal to labels.
      proposal_to_label_map[i] means the index of the ground truth instance for
      the i-th proposal.
    max_num_fg: a integer represents the number of masks per image.
  Returns:
    class_targets, boxes, proposal_to_label_map, box_targets that have
    foreground objects.
  """
  with tf.name_scope('select_fg_for_masks'):
    # Masks are for positive (fg) objects only. Reference: https://github.com/facebookresearch/Detectron/blob/master/detectron/roi_data/mask_rcnn.py  # pylint: disable=line-too-long
    batch_size = boxes.shape[0]
    _, fg_indices = tf.nn.top_k(
        tf.to_float(tf.greater(class_targets, 0)), k=max_num_fg)
    # Contructs indices for gather.
    indices = tf.reshape(
        fg_indices + tf.expand_dims(
            tf.range(batch_size) * tf.shape(class_targets)[1], 1), [-1])

    fg_class_targets = tf.reshape(
        tf.gather(tf.reshape(class_targets, [-1, 1]), indices),
        [batch_size, -1])
    fg_box_targets = tf.reshape(
        tf.gather(tf.reshape(box_targets, [-1, 4]), indices),
        [batch_size, -1, 4])
    fg_box_rois = tf.reshape(
        tf.gather(tf.reshape(boxes, [-1, 4]), indices), [batch_size, -1, 4])
    fg_proposal_to_label_map = tf.reshape(
        tf.gather(tf.reshape(proposal_to_label_map, [-1, 1]), indices),
        [batch_size, -1])

  return (fg_class_targets, fg_box_targets, fg_box_rois,
          fg_proposal_to_label_map)
//  [python]
def select_fg_for_masks(class_targets, box_targets, boxes,
                        proposal_to_label_map, max_num_fg=128):
  """Selects the fore ground objects for mask branch during training.

  Args:
    class_targets: a tensor of shape [batch_size, num_boxes]  representing the
      class label for each box.
    box_targets: a tensor with a shape of [batch_size, num_boxes, 4]. The tensor
      contains the ground truth pixel coordinates of the scaled images for each
      roi.
    boxes: A 3-D Tensor of shape [batch_size, num_boxes, 4]. Each row
      represents a box with [y1, x1, y2, x2] in un-normalized coordinates.
    proposal_to_label_map: a tensor with a shape of [batch_size, num_boxes].
      This tensor keeps the mapping between proposal to labels.
      proposal_to_label_map[i] means the index of the ground truth instance for
      the i-th proposal.
    max_num_fg: a integer represents the number of masks per image.
  Returns:
    class_targets, boxes, proposal_to_label_map, box_targets that have
    foreground objects.
  """
  with tf.name_scope('select_fg_for_masks'):
    # Masks are for positive (fg) objects only. Reference: https://github.com/facebookresearch/Detectron/blob/master/detectron/roi_data/mask_rcnn.py  # pylint: disable=line-too-long
    batch_size = boxes.shape[0]
    _, fg_indices = tf.nn.top_k(
        tf.cast(tf.greater(class_targets, 0), tf.float32), k=max_num_fg)
    # Contructs indices for gather.
    indices = tf.reshape(
        fg_indices + tf.expand_dims(
            tf.range(batch_size) * tf.shape(class_targets)[1], 1), [-1])

    fg_class_targets = tf.reshape(
        tf.gather(tf.reshape(class_targets, [-1, 1]), indices),
        [batch_size, -1])
    fg_box_targets = tf.reshape(
        tf.gather(tf.reshape(box_targets, [-1, 4]), indices),
        [batch_size, -1, 4])
    fg_box_rois = tf.reshape(
        tf.gather(tf.reshape(boxes, [-1, 4]), indices), [batch_size, -1, 4])
    fg_proposal_to_label_map = tf.reshape(
        tf.gather(tf.reshape(proposal_to_label_map, [-1, 1]), indices),
        [batch_size, -1])

  return (fg_class_targets, fg_box_targets, fg_box_rois,
          fg_proposal_to_label_map)
//  [python]
def _create_eval_metric(features, labels, params):
  """Creates eval_metric for model_fn."""
  outputs_to_scales_to_logits = _build_network(
      features, tf.estimator.ModeKeys.EVAL, params)

  semantic_merged_logits = (
      outputs_to_scales_to_logits[common.OUTPUT_TYPE][_MERGED_LOGITS_SCOPE])

  def metric_fn(semantic_merged_logits, labels):
    """Creates metric_fn for TPUEstimatorSpec."""
    logits = tf.image.resize_bilinear(
        semantic_merged_logits, params['crop_size'], align_corners=True)
    predictions_with_shape = tf.argmax(logits, 3, output_type=tf.int32)
    predictions = tf.reshape(predictions_with_shape, shape=[-1])

    labels = tf.reshape(labels, shape=[-1])
    weights = tf.to_float(tf.not_equal(labels, params['ignore_label']))

    # Set ignore_label regions to label 0, because metrics.mean_iou requires
    # range of labels = [0, dataset.num_classes). Note the ignore_lable regions
    # are not evaluated since the corresponding regions contain weights = 0.
    labels = tf.where(
        tf.equal(labels, params['ignore_label']), tf.zeros_like(labels), labels)

    return {
        'miou':
            tf.metrics.mean_iou(
                predictions, labels, params['num_classes'], weights=weights),
    }

  return metric_fn, [semantic_merged_logits, labels]
//  [python]
def _create_eval_metric(features, labels, params):
  """Creates eval_metric for model_fn."""
  outputs_to_scales_to_logits = _build_network(
      features, tf.estimator.ModeKeys.EVAL, params)

  semantic_merged_logits = (
      outputs_to_scales_to_logits[common.OUTPUT_TYPE][_MERGED_LOGITS_SCOPE])

  def metric_fn(semantic_merged_logits, labels):
    """Creates metric_fn for TPUEstimatorSpec."""
    logits = tf.image.resize_bilinear(
        semantic_merged_logits, params['crop_size'], align_corners=True)
    predictions_with_shape = tf.argmax(logits, 3, output_type=tf.int32)
    predictions = tf.reshape(predictions_with_shape, shape=[-1])

    labels = tf.reshape(labels, shape=[-1])
    weights = tf.cast(tf.not_equal(labels, params['ignore_label']), tf.float32)

    # Set ignore_label regions to label 0, because metrics.mean_iou requires
    # range of labels = [0, dataset.num_classes). Note the ignore_lable regions
    # are not evaluated since the corresponding regions contain weights = 0.
    labels = tf.where(
        tf.equal(labels, params['ignore_label']), tf.zeros_like(labels), labels)

    return {
        'miou':
            tf.metrics.mean_iou(
                predictions, labels, params['num_classes'], weights=weights),
    }

  return metric_fn, [semantic_merged_logits, labels]
//  [python]
def _call_cell(self,
                 inputs,
                 initial_cell_state=None,
                 initial_output=None,
                 dtype=None,
                 sequence_length=None):
    """Run this LSTM on inputs, starting from the given state.

    Args:
      inputs: `3-D` tensor with shape `[time_len, batch_size, input_size]`
      initial_cell_state: initial value for cell state, shape `[batch_size,
        self._num_units]`
      initial_output: initial value of cell output, shape `[batch_size,
        self._num_units]`
      dtype: The data type for the initial state and expected output.
      sequence_length: Specifies the length of each sequence in inputs. An
        `int32` or `int64` vector (tensor) size `[batch_size]`, values in `[0,
        time_len)` or None.

    Returns:
      A pair containing:

      - Cell state (cs): A `3-D` tensor of shape `[time_len, batch_size,
                         output_size]`
      - Output (h): A `3-D` tensor of shape `[time_len, batch_size,
                    output_size]`
    """

    inputs_shape = inputs.get_shape().with_rank(3)
    time_len = inputs_shape[0]
    if time_len is None:
      time_len = tf.shape(inputs)[0]

    if self._use_peephole:
      wci = self._w_i_diag
      wco = self._w_o_diag
      wcf = self._w_f_diag
    else:
      wci = wcf = wco = tf.zeros([self._num_units], dtype=dtype)

    if sequence_length is None:
      max_seq_len = tf.to_int64(time_len)
    else:
      max_seq_len = tf.to_int64(tf.reduce_max(sequence_length))

    _, cs, _, _, _, _, h = tf.raw_ops.BlockLSTM(
        seq_len_max=max_seq_len,
        x=inputs,
        cs_prev=initial_cell_state,
        h_prev=initial_output,
        w=self._kernel,
        wci=wci,
        wcf=wcf,
        wco=wco,
        b=self._bias,
        forget_bias=self._forget_bias,
        cell_clip=self._cell_clip,
        use_peephole=self._use_peephole)
    return cs, h
//  [python]
def _call_cell(self,
                 inputs,
                 initial_cell_state=None,
                 initial_output=None,
                 dtype=None,
                 sequence_length=None):
    """Run this LSTM on inputs, starting from the given state.

    Args:
      inputs: `3-D` tensor with shape `[time_len, batch_size, input_size]`
      initial_cell_state: initial value for cell state, shape `[batch_size,
        self._num_units]`
      initial_output: initial value of cell output, shape `[batch_size,
        self._num_units]`
      dtype: The data type for the initial state and expected output.
      sequence_length: Specifies the length of each sequence in inputs. An
        `int32` or `int64` vector (tensor) size `[batch_size]`, values in `[0,
        time_len)` or None.

    Returns:
      A pair containing:

      - Cell state (cs): A `3-D` tensor of shape `[time_len, batch_size,
                         output_size]`
      - Output (h): A `3-D` tensor of shape `[time_len, batch_size,
                    output_size]`
    """

    inputs_shape = inputs.get_shape().with_rank(3)
    time_len = inputs_shape[0]
    if time_len is None:
      time_len = tf.shape(inputs)[0]

    if self._use_peephole:
      wci = self._w_i_diag
      wco = self._w_o_diag
      wcf = self._w_f_diag
    else:
      wci = wcf = wco = tf.zeros([self._num_units], dtype=dtype)

    if sequence_length is None:
      max_seq_len = tf.cast(time_len, tf.int64)
    else:
      max_seq_len = tf.cast(tf.reduce_max(sequence_length), tf.int64)

    _, cs, _, _, _, _, h = tf.raw_ops.BlockLSTM(
        seq_len_max=max_seq_len,
        x=inputs,
        cs_prev=initial_cell_state,
        h_prev=initial_output,
        w=self._kernel,
        wci=wci,
        wcf=wcf,
        wco=wco,
        b=self._bias,
        forget_bias=self._forget_bias,
        cell_clip=self._cell_clip,
        use_peephole=self._use_peephole)
    return cs, h
//  [python]
def _resnet_layer(inputs,
                  num_filters=16,
                  kernel_size=3,
                  strides=1,
                  activation='relu',
                  depth=20,
                  batchnorm=False,
                  conv_first=True,
                  variational=False,
                  n_examples=None):
  """2D Convolution-Batch Normalization-Activation stack builder.

  Args:
    inputs (tensor): input tensor from input image or previous layer
    num_filters (int): Conv2D number of filters
    kernel_size (int): Conv2D square kernel dimensions
    strides (int): Conv2D square stride dimensions
    activation (string): Activation function string.
    depth (int): ResNet depth; used for initialization scale.
    batchnorm (bool): whether to include batch normalization
    conv_first (bool): conv-bn-activation (True) or bn-activation-conv (False)
    variational (bool): Whether to use a variational convolutional layer.
    n_examples (int): Number of examples per epoch for variational KL.

  Returns:
      x (tensor): tensor as input to the next layer
  """
  if variational:

    def fixup_init(shape, dtype=None):
      """Fixup initialization; see https://arxiv.org/abs/1901.09321."""
      return keras.initializers.he_normal()(
          shape, dtype=dtype) * depth**(-1 / 4)

    p_fn, q_fn = mean_field_fn(empirical_bayes=True, initializer=fixup_init)

    def normalized_kl_fn(q, p, _):
      return tfp.distributions.kl_divergence(q, p) / tf.to_float(n_examples)

    conv = tfp.layers.Convolution2DFlipout(
        num_filters,
        kernel_size=kernel_size,
        strides=strides,
        padding='same',
        kernel_prior_fn=p_fn,
        kernel_posterior_fn=q_fn,
        kernel_divergence_fn=normalized_kl_fn)
  else:
    conv = keras.layers.Conv2D(
        num_filters,
        kernel_size=kernel_size,
        strides=strides,
        padding='same',
        kernel_initializer='he_normal',
        kernel_regularizer=keras.regularizers.l2(1e-4))

  def apply_conv(net):
    return conv(net)

  x = inputs
  x = apply_conv(x) if conv_first else x
  if batchnorm:
    x = keras.layers.BatchNormalization()(x)
  x = keras.layers.Activation(activation)(x) if activation is not None else x
  x = x if conv_first else apply_conv(x)
  return x
//  [python]
def _resnet_layer(inputs,
                  num_filters=16,
                  kernel_size=3,
                  strides=1,
                  activation='relu',
                  depth=20,
                  batchnorm=False,
                  conv_first=True,
                  variational=False,
                  n_examples=None):
  """2D Convolution-Batch Normalization-Activation stack builder.

  Args:
    inputs (tensor): input tensor from input image or previous layer
    num_filters (int): Conv2D number of filters
    kernel_size (int): Conv2D square kernel dimensions
    strides (int): Conv2D square stride dimensions
    activation (string): Activation function string.
    depth (int): ResNet depth; used for initialization scale.
    batchnorm (bool): whether to include batch normalization
    conv_first (bool): conv-bn-activation (True) or bn-activation-conv (False)
    variational (bool): Whether to use a variational convolutional layer.
    n_examples (int): Number of examples per epoch for variational KL.

  Returns:
      x (tensor): tensor as input to the next layer
  """
  if variational:

    def fixup_init(shape, dtype=None):
      """Fixup initialization; see https://arxiv.org/abs/1901.09321."""
      return keras.initializers.he_normal()(
          shape, dtype=dtype) * depth**(-1 / 4)

    p_fn, q_fn = mean_field_fn(empirical_bayes=True, initializer=fixup_init)

    def normalized_kl_fn(q, p, _):
      return tfp.distributions.kl_divergence(q, p) / tf.cast(n_examples, tf.float32)

    conv = tfp.layers.Convolution2DFlipout(
        num_filters,
        kernel_size=kernel_size,
        strides=strides,
        padding='same',
        kernel_prior_fn=p_fn,
        kernel_posterior_fn=q_fn,
        kernel_divergence_fn=normalized_kl_fn)
  else:
    conv = keras.layers.Conv2D(
        num_filters,
        kernel_size=kernel_size,
        strides=strides,
        padding='same',
        kernel_initializer='he_normal',
        kernel_regularizer=keras.regularizers.l2(1e-4))

  def apply_conv(net):
    return conv(net)

  x = inputs
  x = apply_conv(x) if conv_first else x
  if batchnorm:
    x = keras.layers.BatchNormalization()(x)
  x = keras.layers.Activation(activation)(x) if activation is not None else x
  x = x if conv_first else apply_conv(x)
  return x
//  [python]
def rpn_loss(score_outputs, box_outputs, labels, params):
  """Computes total RPN detection loss.

  Computes total RPN detection loss including box and score from all levels.
  Args:
    score_outputs: an OrderDict with keys representing levels and values
      representing scores in [batch_size, height, width, num_anchors].
    box_outputs: an OrderDict with keys representing levels and values
      representing box regression targets in
      [batch_size, height, width, num_anchors * 4].
    labels: the dictionary that returned from dataloader that includes
      groundturth targets.
    params: the dictionary including training parameters specified in
      default_haprams() in mask_rcnn_params.py.
  Returns:
    total_rpn_loss: a float tensor representing total loss reduced from
      score and box losses from all levels.
    rpn_score_loss: a float tensor representing total score loss.
    rpn_box_loss: a float tensor representing total box regression loss.
  """
  with tf.name_scope('rpn_loss'):
    levels = score_outputs.keys()

    score_losses = []
    box_losses = []
    for level in levels:
      b, h, w, n = score_outputs[level].get_shape().as_list()
      score_targets_at_level = tf.reshape(
          labels['score_targets_%d' % level], [b, h, w, n])
      box_targets_at_level = tf.reshape(
          labels['box_targets_%d' % level], [b, h, w, n * 4])
      score_losses.append(
          _rpn_score_loss(
              score_outputs[level],
              score_targets_at_level,
              normalizer=tf.to_float(b * params['rpn_batch_size_per_im'])))
      box_losses.append(
          _rpn_box_loss(box_outputs[level], box_targets_at_level))

    # Sum per level losses to total loss.
    rpn_score_loss = tf.add_n(score_losses)
    rpn_box_loss = params['rpn_box_loss_weight'] * tf.add_n(box_losses)
    total_rpn_loss = rpn_score_loss + rpn_box_loss
    return total_rpn_loss, rpn_score_loss, rpn_box_loss
//  [python]
def rpn_loss(score_outputs, box_outputs, labels, params):
  """Computes total RPN detection loss.

  Computes total RPN detection loss including box and score from all levels.
  Args:
    score_outputs: an OrderDict with keys representing levels and values
      representing scores in [batch_size, height, width, num_anchors].
    box_outputs: an OrderDict with keys representing levels and values
      representing box regression targets in
      [batch_size, height, width, num_anchors * 4].
    labels: the dictionary that returned from dataloader that includes
      groundturth targets.
    params: the dictionary including training parameters specified in
      default_haprams() in mask_rcnn_params.py.
  Returns:
    total_rpn_loss: a float tensor representing total loss reduced from
      score and box losses from all levels.
    rpn_score_loss: a float tensor representing total score loss.
    rpn_box_loss: a float tensor representing total box regression loss.
  """
  with tf.name_scope('rpn_loss'):
    levels = score_outputs.keys()

    score_losses = []
    box_losses = []
    for level in levels:
      b, h, w, n = score_outputs[level].get_shape().as_list()
      score_targets_at_level = tf.reshape(
          labels['score_targets_%d' % level], [b, h, w, n])
      box_targets_at_level = tf.reshape(
          labels['box_targets_%d' % level], [b, h, w, n * 4])
      score_losses.append(
          _rpn_score_loss(
              score_outputs[level],
              score_targets_at_level,
              normalizer=tf.cast(b * params['rpn_batch_size_per_im'], tf.float32)))
      box_losses.append(
          _rpn_box_loss(box_outputs[level], box_targets_at_level))

    # Sum per level losses to total loss.
    rpn_score_loss = tf.add_n(score_losses)
    rpn_box_loss = params['rpn_box_loss_weight'] * tf.add_n(box_losses)
    total_rpn_loss = rpn_score_loss + rpn_box_loss
    return total_rpn_loss, rpn_score_loss, rpn_box_loss
//  [python]
def sigmoid_rampup(global_step, rampup_length):
    global_step = tf.to_float(global_step)
    rampup_length = tf.to_float(rampup_length)
    def ramp():
        phase = 1.0 - tf.maximum(0.0, global_step) / rampup_length
        return tf.exp(-5.0 * phase * phase)

    result = tf.cond(global_step < rampup_length, ramp, lambda: tf.constant(1.0))
    return tf.identity(result, name="sigmoid_rampup")
//  [python]
def sigmoid_rampup(global_step, rampup_length):
    global_step = tf.cast(global_step, tf.float32)
    rampup_length = tf.cast(rampup_length, tf.float32)
    def ramp():
        phase = 1.0 - tf.maximum(0.0, global_step) / rampup_length
        return tf.exp(-5.0 * phase * phase)

    result = tf.cond(global_step < rampup_length, ramp, lambda: tf.constant(1.0))
    return tf.identity(result, name="sigmoid_rampup")
//  [python]
def sigmoid_rampdown(global_step, rampdown_length, training_length):
    global_step = tf.to_float(global_step)
    rampdown_length = tf.to_float(rampdown_length)
    training_length = tf.to_float(training_length)
    def ramp():
        phase = 1.0 - tf.maximum(0.0, training_length - global_step) / rampdown_length
        return tf.exp(-12.5 * phase * phase)

    result = tf.cond(global_step >= training_length - rampdown_length,
                     ramp,
                     lambda: tf.constant(1.0))
    return tf.identity(result, name="sigmoid_rampdown")
//  [python]
def sigmoid_rampdown(global_step, rampdown_length, training_length):
    global_step = tf.cast(global_step, tf.float32)
    rampdown_length = tf.cast(rampdown_length, tf.float32)
    training_length = tf.cast(training_length, tf.float32)
    def ramp():
        phase = 1.0 - tf.maximum(0.0, training_length - global_step) / rampdown_length
        return tf.exp(-12.5 * phase * phase)

    result = tf.cond(global_step >= training_length - rampdown_length,
                     ramp,
                     lambda: tf.constant(1.0))
    return tf.identity(result, name="sigmoid_rampdown")
//  [python]
def _create_detection_op(model, input_dict, batch):
  """Create detection ops.

  Args:
    model: model to perform predictions with.
    input_dict: A dict holds input data.
    batch: batch size for evaluation.

  Returns:
    Detection tensor ops.
  """
  video_tensor = tf.stack(list(input_dict[fields.InputDataFields.image]))
  preprocessed_video, true_image_shapes = model.preprocess(
      tf.to_float(video_tensor))
  if batch is not None:
    prediction_dict = model.predict(preprocessed_video, true_image_shapes,
                                    batch)
  else:
    prediction_dict = model.predict(preprocessed_video, true_image_shapes)

  return model.postprocess(prediction_dict, true_image_shapes)
//  [python]
def _create_detection_op(model, input_dict, batch):
  """Create detection ops.

  Args:
    model: model to perform predictions with.
    input_dict: A dict holds input data.
    batch: batch size for evaluation.

  Returns:
    Detection tensor ops.
  """
  video_tensor = tf.stack(list(input_dict[fields.InputDataFields.image]))
  preprocessed_video, true_image_shapes = model.preprocess(
      tf.cast(video_tensor, tf.float32))
  if batch is not None:
    prediction_dict = model.predict(preprocessed_video, true_image_shapes,
                                    batch)
  else:
    prediction_dict = model.predict(preprocessed_video, true_image_shapes)

  return model.postprocess(prediction_dict, true_image_shapes)
//  [python]
def classification_costs(logits, labels, name=None):
    """Compute classification cost mean and classification cost per sample

    Assume unlabeled examples have label == -1. For unlabeled examples, cost == 0.
    Compute the mean over all examples.
    Note that unlabeled examples are treated differently in error calculation.
    """
    with tf.name_scope(name, "classification_costs") as scope:
        applicable = tf.not_equal(labels, -1)

        # Change -1s to zeros to make cross-entropy computable
        labels = tf.where(applicable, labels, tf.zeros_like(labels))

        # This will now have incorrect values for unlabeled examples
        per_sample = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits, labels=labels)

        # Retain costs only for labeled
        per_sample = tf.where(applicable, per_sample, tf.zeros_like(per_sample))

        # Take mean over all examples, not just labeled examples.
        labeled_sum = tf.reduce_sum(per_sample)
        total_count = tf.to_float(tf.shape(per_sample)[0])
        mean = tf.div(labeled_sum, total_count, name=scope)

        return mean, per_sample
//  [python]
def classification_costs(logits, labels, name=None):
    """Compute classification cost mean and classification cost per sample

    Assume unlabeled examples have label == -1. For unlabeled examples, cost == 0.
    Compute the mean over all examples.
    Note that unlabeled examples are treated differently in error calculation.
    """
    with tf.name_scope(name, "classification_costs") as scope:
        applicable = tf.not_equal(labels, -1)

        # Change -1s to zeros to make cross-entropy computable
        labels = tf.where(applicable, labels, tf.zeros_like(labels))

        # This will now have incorrect values for unlabeled examples
        per_sample = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits, labels=labels)

        # Retain costs only for labeled
        per_sample = tf.where(applicable, per_sample, tf.zeros_like(per_sample))

        # Take mean over all examples, not just labeled examples.
        labeled_sum = tf.reduce_sum(per_sample)
        total_count = tf.cast(tf.shape(per_sample)[0], tf.float32)
        mean = tf.div(labeled_sum, total_count, name=scope)

        return mean, per_sample
//  [python]
def _build_lm_task(cfg, dataset):
  """Builds a language modeling task."""

  def _build(batch):
    """Builds the sonnet module."""
    # Shape is [batch size, sequence length]
    inp = batch["text"]

    # Clip the vocab to be at most vocab_size.
    inp = tf.minimum(inp,
                     tf.to_int64(tf.reshape(cfg["vocab_size"] - 1, [1, 1])))

    embed = snt.Embed(vocab_size=cfg["vocab_size"], embed_dim=cfg["embed_dim"])
    embedded_chars = embed(inp)

    rnn = utils.get_rnn_core(cfg["core"])
    batch_size = inp.shape.as_list()[0]

    state = rnn.initial_state(batch_size, trainable=cfg["trainable_init"])

    outputs, state = tf.nn.dynamic_rnn(rnn, embedded_chars, initial_state=state)

    w_init = utils.get_initializer(cfg["w_init"])
    pred_logits = snt.BatchApply(
        snt.Linear(cfg["vocab_size"], initializers={"w": w_init}))(
            outputs[:, :-1])
    actual_output_tokens = inp[:, 1:]

    flat_s = [pred_logits.shape[0] * pred_logits.shape[1], pred_logits.shape[2]]
    flat_pred_logits = tf.reshape(pred_logits, flat_s)
    flat_actual_tokens = tf.reshape(actual_output_tokens, [flat_s[0]])

    loss_vec = tf.nn.sparse_softmax_cross_entropy_with_logits(
        labels=flat_actual_tokens, logits=flat_pred_logits)
    return tf.reduce_mean(loss_vec)

  return base.DatasetModelTask(lambda: snt.Module(_build), dataset)
//  [python]
def _build_lm_task(cfg, dataset):
  """Builds a language modeling task."""

  def _build(batch):
    """Builds the sonnet module."""
    # Shape is [batch size, sequence length]
    inp = batch["text"]

    # Clip the vocab to be at most vocab_size.
    inp = tf.minimum(inp,
                     tf.cast(tf.reshape(cfg["vocab_size"] - 1, [1, 1]), tf.int64))

    embed = snt.Embed(vocab_size=cfg["vocab_size"], embed_dim=cfg["embed_dim"])
    embedded_chars = embed(inp)

    rnn = utils.get_rnn_core(cfg["core"])
    batch_size = inp.shape.as_list()[0]

    state = rnn.initial_state(batch_size, trainable=cfg["trainable_init"])

    outputs, state = tf.nn.dynamic_rnn(rnn, embedded_chars, initial_state=state)

    w_init = utils.get_initializer(cfg["w_init"])
    pred_logits = snt.BatchApply(
        snt.Linear(cfg["vocab_size"], initializers={"w": w_init}))(
            outputs[:, :-1])
    actual_output_tokens = inp[:, 1:]

    flat_s = [pred_logits.shape[0] * pred_logits.shape[1], pred_logits.shape[2]]
    flat_pred_logits = tf.reshape(pred_logits, flat_s)
    flat_actual_tokens = tf.reshape(actual_output_tokens, [flat_s[0]])

    loss_vec = tf.nn.sparse_softmax_cross_entropy_with_logits(
        labels=flat_actual_tokens, logits=flat_pred_logits)
    return tf.reduce_mean(loss_vec)

  return base.DatasetModelTask(lambda: snt.Module(_build), dataset)
//  [python]
def encode(self, x, noise):
    x = tf.to_float(x)
    # we can't use tf.pow(..., 8.0) because of a high-error approximation
    # on TPU.  Instead we square three times.
    x = tf.sign(x) * tf.square(tf.square(tf.square(tf.abs(x) * 128.0)))
    x = _to_bfloat16_unbiased(x, noise)
    return x
//  [python]
def encode(self, x, noise):
    x = tf.cast(x, tf.float32)
    # we can't use tf.pow(..., 8.0) because of a high-error approximation
    # on TPU.  Instead we square three times.
    x = tf.sign(x) * tf.square(tf.square(tf.square(tf.abs(x) * 128.0)))
    x = _to_bfloat16_unbiased(x, noise)
    return x
//  [python]
def decode(self, x):
    x = tf.to_float(x)
    # we can't use tf.pow(..., 0.125) because of a high-error approximation
    # on TPU.  Instead we sqrt three times.
    return tf.sign(x) * (tf.sqrt(tf.sqrt(tf.sqrt(tf.abs(x)))) / 128.0)
//  [python]
def decode(self, x):
    x = tf.cast(x, tf.float32)
    # we can't use tf.pow(..., 0.125) because of a high-error approximation
    # on TPU.  Instead we sqrt three times.
    return tf.sign(x) * (tf.sqrt(tf.sqrt(tf.sqrt(tf.abs(x)))) / 128.0)
//  [python]
def noise_from_step_num():
  """Quantization noise equal to (phi * (step_num + 1)) mod 1.0.
  Not using random_uniform here due to a problem on TPU in that random seeds
  are not respected, which may cause the parameters on different replicas
  to go out-of-sync.
  Returns:
    a float32 scalar
  """
  step = tf.to_int32(tf.train.get_or_create_global_step()) + 1
  phi = ((5 ** 0.5) - 1) / 2
  # Naive computation tf.mod(phi * step, 1.0) in float32 would be disastrous
  # due to loss of precision when the step number gets large.
  # Computation in doubles does not work on TPU, so we use this complicated
  # alternative computation which does not suffer from these roundoff errors.
  ret = 0.0
  for i in range(30):
    ret += (((phi * (2 ** i)) % 1.0)  # double-precision computation in python
            * tf.to_float(tf.mod(step // (2 ** i), 2)))
  return tf.mod(ret, 1.0)
//  [python]
def noise_from_step_num():
  """Quantization noise equal to (phi * (step_num + 1)) mod 1.0.
  Not using random_uniform here due to a problem on TPU in that random seeds
  are not respected, which may cause the parameters on different replicas
  to go out-of-sync.
  Returns:
    a float32 scalar
  """
  step = tf.cast(tf.train.get_or_create_global_step(), tf.int32) + 1
  phi = ((5 ** 0.5) - 1) / 2
  # Naive computation tf.mod(phi * step, 1.0) in float32 would be disastrous
  # due to loss of precision when the step number gets large.
  # Computation in doubles does not work on TPU, so we use this complicated
  # alternative computation which does not suffer from these roundoff errors.
  ret = 0.0
  for i in range(30):
    ret += (((phi * (2 ** i)) % 1.0)  # double-precision computation in python
            * tf.to_float(tf.mod(step // (2 ** i), 2)))
  return tf.mod(ret, 1.0)
//  [python]
def _randomized_roundoff_to_bfloat16(x, noise, cand1, cand2):
  """Round-off x to cand1 or to cand2 in an unbiased way.
  Cand1 and cand2 are the same shape as x.
  For every element of x, the corresponding elements of cand1 and cand2 should
  be the two closest bfloat16 values to x.  Order does not matter.
  cand1 and cand2 must differ from each other.
  Args:
    x: A float32 Tensor.
    noise: A Tensor broadcastable to the shape of x containing
    random uniform values in [0.0, 1.0].
    cand1: A bfloat16 Tensor the same shape as x.
    cand2: A bfloat16 Tensor the same shape as x.
  Returns:
    A bfloat16 Tensor.
  """
  cand1_f = tf.to_float(cand1)
  cand2_f = tf.to_float(cand2)
  step_size = cand2_f - cand1_f
  fpart = (x - cand1_f) / step_size
  ret = tf.where(tf.greater(fpart, noise), cand2, cand1)
  return ret
//  [python]
def _randomized_roundoff_to_bfloat16(x, noise, cand1, cand2):
  """Round-off x to cand1 or to cand2 in an unbiased way.
  Cand1 and cand2 are the same shape as x.
  For every element of x, the corresponding elements of cand1 and cand2 should
  be the two closest bfloat16 values to x.  Order does not matter.
  cand1 and cand2 must differ from each other.
  Args:
    x: A float32 Tensor.
    noise: A Tensor broadcastable to the shape of x containing
    random uniform values in [0.0, 1.0].
    cand1: A bfloat16 Tensor the same shape as x.
    cand2: A bfloat16 Tensor the same shape as x.
  Returns:
    A bfloat16 Tensor.
  """
  cand1_f = tf.cast(cand1, tf.float32)
  cand2_f = tf.cast(cand2, tf.float32)
  step_size = cand2_f - cand1_f
  fpart = (x - cand1_f) / step_size
  ret = tf.where(tf.greater(fpart, noise), cand2, cand1)
  return ret
//  [python]
def _to_bfloat16_unbiased(x, noise):
  """Convert a float32 to a bfloat16 using randomized roundoff.
  Args:
    x: A float32 Tensor.
    noise: a float32 Tensor with values in [0, 1), broadcastable to tf.shape(x)
  Returns:
    A float32 Tensor.
  """
  x_sign = tf.sign(x)
  # Make sure x is positive.  If it is zero, the two candidates are identical.
  x = x * x_sign + 1e-30
  cand1 = tf.to_bfloat16(x)
  cand1_f = tf.to_float(cand1)
  # This relies on the fact that for a positive bfloat16 b,
  # b * 1.005 gives you the next higher bfloat16 and b*0.995 gives you the
  # next lower one. Both 1.005 and 0.995 are ballpark estimation.
  cand2 = tf.to_bfloat16(
      tf.where(tf.greater(x, cand1_f), cand1_f * 1.005, cand1_f * 0.995))
  ret = _randomized_roundoff_to_bfloat16(x, noise, cand1, cand2)
  return ret * tf.to_bfloat16(x_sign)
//  [python]
def _to_bfloat16_unbiased(x, noise):
  """Convert a float32 to a bfloat16 using randomized roundoff.
  Args:
    x: A float32 Tensor.
    noise: a float32 Tensor with values in [0, 1), broadcastable to tf.shape(x)
  Returns:
    A float32 Tensor.
  """
  x_sign = tf.sign(x)
  # Make sure x is positive.  If it is zero, the two candidates are identical.
  x = x * x_sign + 1e-30
  cand1 = tf.cast(x, tf.bfloat16)
  cand1_f = tf.cast(cand1, tf.float32)
  # This relies on the fact that for a positive bfloat16 b,
  # b * 1.005 gives you the next higher bfloat16 and b*0.995 gives you the
  # next lower one. Both 1.005 and 0.995 are ballpark estimation.
  cand2 = tf.cast(
      tf.where(tf.greater(x, cand1_f), cand1_f * 1.005, cand1_f * 0.995), tf.bfloat16)
  ret = _randomized_roundoff_to_bfloat16(x, noise, cand1, cand2)
  return ret * tf.cast(x_sign, tf.bfloat16)
//  [python]
def preprocess_for_eval(image,
                        output_height,
                        output_width,
                        resize_side,
                        use_grayscale=False):
  """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.
    use_grayscale: Whether to convert the image from RGB to grayscale.

  Returns:
    A preprocessed image.
  """
  image = _aspect_preserving_resize(image, resize_side)
  image = _central_crop([image], output_height, output_width)[0]
  image.set_shape([output_height, output_width, 3])
  image = tf.to_float(image)
  if use_grayscale:
    image = tf.image.rgb_to_grayscale(image)
  return _mean_image_subtraction(image, [_R_MEAN, _G_MEAN, _B_MEAN])
//  [python]
def preprocess_for_eval(image,
                        output_height,
                        output_width,
                        resize_side,
                        use_grayscale=False):
  """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.
    use_grayscale: Whether to convert the image from RGB to grayscale.

  Returns:
    A preprocessed image.
  """
  image = _aspect_preserving_resize(image, resize_side)
  image = _central_crop([image], output_height, output_width)[0]
  image.set_shape([output_height, output_width, 3])
  image = tf.cast(image, tf.float32)
  if use_grayscale:
    image = tf.image.rgb_to_grayscale(image)
  return _mean_image_subtraction(image, [_R_MEAN, _G_MEAN, _B_MEAN])
//  [python]
def _crop(image, offset_height, offset_width, crop_height, crop_width):
  """Crops the given image using the provided offsets and sizes.

  Note that the method doesn't assume we know the input image size but it does
  assume we know the input image rank.

  Args:
    image: an image of shape [height, width, channels].
    offset_height: a scalar tensor indicating the height offset.
    offset_width: a scalar tensor indicating the width offset.
    crop_height: the height of the cropped image.
    crop_width: the width of the cropped image.

  Returns:
    the cropped (and resized) image.

  Raises:
    InvalidArgumentError: if the rank is not 3 or if the image dimensions are
      less than the crop size.
  """
  original_shape = tf.shape(image)

  rank_assertion = tf.Assert(
      tf.equal(tf.rank(image), 3),
      ['Rank of image must be equal to 3.'])
  with tf.control_dependencies([rank_assertion]):
    cropped_shape = tf.stack([crop_height, crop_width, original_shape[2]])

  size_assertion = tf.Assert(
      tf.logical_and(
          tf.greater_equal(original_shape[0], crop_height),
          tf.greater_equal(original_shape[1], crop_width)),
      ['Crop size greater than the image size.'])

  offsets = tf.to_int32(tf.stack([offset_height, offset_width, 0]))

  # Use tf.slice instead of crop_to_bounding box as it accepts tensors to
  # define the crop size.
  with tf.control_dependencies([size_assertion]):
    image = tf.slice(image, offsets, cropped_shape)
  return tf.reshape(image, cropped_shape)
//  [python]
def _crop(image, offset_height, offset_width, crop_height, crop_width):
  """Crops the given image using the provided offsets and sizes.

  Note that the method doesn't assume we know the input image size but it does
  assume we know the input image rank.

  Args:
    image: an image of shape [height, width, channels].
    offset_height: a scalar tensor indicating the height offset.
    offset_width: a scalar tensor indicating the width offset.
    crop_height: the height of the cropped image.
    crop_width: the width of the cropped image.

  Returns:
    the cropped (and resized) image.

  Raises:
    InvalidArgumentError: if the rank is not 3 or if the image dimensions are
      less than the crop size.
  """
  original_shape = tf.shape(image)

  rank_assertion = tf.Assert(
      tf.equal(tf.rank(image), 3),
      ['Rank of image must be equal to 3.'])
  with tf.control_dependencies([rank_assertion]):
    cropped_shape = tf.stack([crop_height, crop_width, original_shape[2]])

  size_assertion = tf.Assert(
      tf.logical_and(
          tf.greater_equal(original_shape[0], crop_height),
          tf.greater_equal(original_shape[1], crop_width)),
      ['Crop size greater than the image size.'])

  offsets = tf.cast(tf.stack([offset_height, offset_width, 0]), tf.int32)

  # Use tf.slice instead of crop_to_bounding box as it accepts tensors to
  # define the crop size.
  with tf.control_dependencies([size_assertion]):
    image = tf.slice(image, offsets, cropped_shape)
  return tf.reshape(image, cropped_shape)
//  [python]
def _smallest_size_at_least(height, width, smallest_side):
  """Computes new shape with the smallest side equal to `smallest_side`.

  Computes new shape with the smallest side equal to `smallest_side` while
  preserving the original aspect ratio.

  Args:
    height: an int32 scalar tensor indicating the current height.
    width: an int32 scalar tensor indicating the current width.
    smallest_side: A python integer or scalar `Tensor` indicating the size of
      the smallest side after resize.

  Returns:
    new_height: an int32 scalar tensor indicating the new height.
    new_width: and int32 scalar tensor indicating the new width.
  """
  smallest_side = tf.convert_to_tensor(smallest_side, dtype=tf.int32)

  height = tf.to_float(height)
  width = tf.to_float(width)
  smallest_side = tf.to_float(smallest_side)

  scale = tf.cond(tf.greater(height, width),
                  lambda: smallest_side / width,
                  lambda: smallest_side / height)
  new_height = tf.to_int32(tf.rint(height * scale))
  new_width = tf.to_int32(tf.rint(width * scale))
  return new_height, new_width
//  [python]
def _smallest_size_at_least(height, width, smallest_side):
  """Computes new shape with the smallest side equal to `smallest_side`.

  Computes new shape with the smallest side equal to `smallest_side` while
  preserving the original aspect ratio.

  Args:
    height: an int32 scalar tensor indicating the current height.
    width: an int32 scalar tensor indicating the current width.
    smallest_side: A python integer or scalar `Tensor` indicating the size of
      the smallest side after resize.

  Returns:
    new_height: an int32 scalar tensor indicating the new height.
    new_width: and int32 scalar tensor indicating the new width.
  """
  smallest_side = tf.convert_to_tensor(smallest_side, dtype=tf.int32)

  height = tf.cast(height, tf.float32)
  width = tf.cast(width, tf.float32)
  smallest_side = tf.cast(smallest_side, tf.float32)

  scale = tf.cond(tf.greater(height, width),
                  lambda: smallest_side / width,
                  lambda: smallest_side / height)
  new_height = tf.to_int32(tf.rint(height * scale))
  new_width = tf.to_int32(tf.rint(width * scale))
  return new_height, new_width
//  [python]
def rpn_loss(score_outputs, box_outputs, labels, params):
  """Computes total RPN detection loss.

  Computes total RPN detection loss including box and score from all levels.
  Args:
    score_outputs: an OrderDict with keys representing levels and values
      representing scores in [batch_size, height, width, num_anchors].
    box_outputs: an OrderDict with keys representing levels and values
      representing box regression targets in
      [batch_size, height, width, num_anchors * 4].
    labels: the dictionary that returned from dataloader that includes
      groundturth targets.
    params: the dictionary including training parameters specified in
      default_haprams function in this file.
  Returns:
    total_rpn_loss: a float tensor representing total loss reduced from
      score and box losses from all levels.
    rpn_score_loss: a float tensor representing total score loss.
    rpn_box_loss: a float tensor representing total box regression loss.
  """
  with tf.name_scope('rpn_loss'):
    levels = score_outputs.keys()

    score_losses = []
    box_losses = []
    for level in levels:
      score_targets_at_level = labels['score_targets_%d' % level]
      box_targets_at_level = labels['box_targets_%d' % level]
      score_losses.append(
          _rpn_score_loss(
              score_outputs[level],
              score_targets_at_level,
              normalizer=tf.to_float(
                  params['batch_size'] * params['rpn_batch_size_per_im'])))
      box_losses.append(
          _rpn_box_loss(box_outputs[level], box_targets_at_level))

    # Sum per level losses to total loss.
    rpn_score_loss = tf.add_n(score_losses)
    rpn_box_loss = params['rpn_box_loss_weight'] * tf.add_n(box_losses)
    total_rpn_loss = rpn_score_loss + rpn_box_loss
    return total_rpn_loss, rpn_score_loss, rpn_box_loss
//  [python]
def rpn_loss(score_outputs, box_outputs, labels, params):
  """Computes total RPN detection loss.

  Computes total RPN detection loss including box and score from all levels.
  Args:
    score_outputs: an OrderDict with keys representing levels and values
      representing scores in [batch_size, height, width, num_anchors].
    box_outputs: an OrderDict with keys representing levels and values
      representing box regression targets in
      [batch_size, height, width, num_anchors * 4].
    labels: the dictionary that returned from dataloader that includes
      groundturth targets.
    params: the dictionary including training parameters specified in
      default_haprams function in this file.
  Returns:
    total_rpn_loss: a float tensor representing total loss reduced from
      score and box losses from all levels.
    rpn_score_loss: a float tensor representing total score loss.
    rpn_box_loss: a float tensor representing total box regression loss.
  """
  with tf.name_scope('rpn_loss'):
    levels = score_outputs.keys()

    score_losses = []
    box_losses = []
    for level in levels:
      score_targets_at_level = labels['score_targets_%d' % level]
      box_targets_at_level = labels['box_targets_%d' % level]
      score_losses.append(
          _rpn_score_loss(
              score_outputs[level],
              score_targets_at_level,
              normalizer=tf.cast(
                  params['batch_size'] * params['rpn_batch_size_per_im'], tf.float32)))
      box_losses.append(
          _rpn_box_loss(box_outputs[level], box_targets_at_level))

    # Sum per level losses to total loss.
    rpn_score_loss = tf.add_n(score_losses)
    rpn_box_loss = params['rpn_box_loss_weight'] * tf.add_n(box_losses)
    total_rpn_loss = rpn_score_loss + rpn_box_loss
    return total_rpn_loss, rpn_score_loss, rpn_box_loss
//  [python]
def preprocess_for_train(image,
                         output_height,
                         output_width,
                         resize_side_min=_RESIZE_SIDE_MIN,
                         resize_side_max=_RESIZE_SIDE_MAX,
                         use_grayscale=False):
  """Preprocesses the given image for training.

  Note that the actual resizing scale is sampled from
    [`resize_size_min`, `resize_size_max`].

  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_min: The lower bound for the smallest side of the image for
      aspect-preserving resizing.
    resize_side_max: The upper bound for the smallest side of the image for
      aspect-preserving resizing.
    use_grayscale: Whether to convert the image from RGB to grayscale.

  Returns:
    A preprocessed image.
  """
  resize_side = tf.random_uniform(
      [], minval=resize_side_min, maxval=resize_side_max+1, dtype=tf.int32)

  image = _aspect_preserving_resize(image, resize_side)
  image = _random_crop([image], output_height, output_width)[0]
  image.set_shape([output_height, output_width, 3])
  image = tf.to_float(image)
  if use_grayscale:
    image = tf.image.rgb_to_grayscale(image)
  image = tf.image.random_flip_left_right(image)
  return _mean_image_subtraction(image, [_R_MEAN, _G_MEAN, _B_MEAN])