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

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

                
                siamese_loss = self.l2_loss
                #if self.threshold is not None:
                if False:
                    ind = tf.unstack(siamese_loss)
                    siamese_loss = [ tf.cond(tf.greater(x, self.threshold), 
                                            lambda: self.threshold + self.threshold * tf.log(x / self.threshold),
                                            lambda: x) for x in ind ]
                    self.robust_l2_loss = siamese_loss
                    siamese_loss = tf.stack(siamese_loss)
                          
                self.siamese_loss = tf.reduce_sum(siamese_loss) / self.cfg['batch_size'] 
        tf.add_to_collection(tf.GraphKeys.LOSSES, self.siamese_loss)

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

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

                
                siamese_loss = self.l2_loss
                #if self.threshold is not None:
                if False:
                    ind = tf.unstack(siamese_loss)
                    siamese_loss = [ tf.cond(tf.greater(x, self.threshold), 
                                            lambda: self.threshold + self.threshold * tf.log(x / self.threshold),
                                            lambda: x) for x in ind ]
                    self.robust_l2_loss = siamese_loss
                    siamese_loss = tf.stack(siamese_loss)
                          
                self.siamese_loss = tf.reduce_sum(siamese_loss) / self.cfg['batch_size'] 
        tf.add_to_collection(tf.GraphKeys.LOSSES, self.siamese_loss)

        losses = [self.siamese_loss]
        return losses
//  [python]
def mpi_resample_cube(mpi, tgt, intrinsics, depth_planes, side_length,
                      cube_res):
  """Resample MPI onto cube centered at target point.

  Args:
    mpi: [B,H,W,D,C], input MPI
    tgt: [B,3], [x,y,z] coordinates for cube center (in reference/mpi frame)
    intrinsics: [B,3,3], MPI reference camera intrinsics
    depth_planes: [D] depth values for MPI planes
    side_length: metric side length of cube
    cube_res: resolution of each cube dimension

  Returns:
    resampled: [B, cube_res, cube_res, cube_res, C]
  """

  batch_size = tf.shape(mpi)[0]
  num_depths = tf.shape(mpi)[3]

  # compute MPI world coordinates
  intrinsics_tile = tf.tile(intrinsics, [num_depths, 1, 1])

  # create cube coordinates
  b_vals = tf.to_float(tf.range(batch_size))
  x_vals = tf.linspace(-side_length / 2.0, side_length / 2.0, cube_res)
  y_vals = tf.linspace(-side_length / 2.0, side_length / 2.0, cube_res)
  z_vals = tf.linspace(side_length / 2.0, -side_length / 2.0, cube_res)
  b, y, x, z = tf.meshgrid(b_vals, y_vals, x_vals, z_vals, indexing='ij')

  x = x + tgt[:, 0, tf.newaxis, tf.newaxis, tf.newaxis]
  y = y + tgt[:, 1, tf.newaxis, tf.newaxis, tf.newaxis]
  z = z + tgt[:, 2, tf.newaxis, tf.newaxis, tf.newaxis]

  ones = tf.ones_like(x)
  coords = tf.stack([x, y, z, ones], axis=1)
  coords_r = tf.reshape(
      tf.transpose(coords, [0, 4, 1, 2, 3]),
      [batch_size * cube_res, 4, cube_res, cube_res])

  # store elements with negative z vals for projection
  bad_inds = tf.less(z, 0.0)

  # project into reference camera to transform coordinates into MPI indices
  filler = tf.constant([0.0, 0.0, 0.0, 1.0], shape=[1, 1, 4])
  filler = tf.tile(filler, [batch_size * cube_res, 1, 1])
  intrinsics_tile = tf.tile(intrinsics, [cube_res, 1, 1])
  intrinsics_tile_4 = tf.concat(
      [intrinsics_tile,
       tf.zeros([batch_size * cube_res, 3, 1])], axis=2)
  intrinsics_tile_4 = tf.concat([intrinsics_tile_4, filler], axis=1)
  coords_proj = cam2pixel(coords_r, intrinsics_tile_4)
  coords_depths = tf.transpose(coords_r[:, 2:3, :, :], [0, 2, 3, 1])
  coords_depth_inds = (tf.to_float(num_depths) - 1) * (
      (1.0 / coords_depths) -
      (1.0 / depth_planes[0])) / ((1.0 / depth_planes[-1]) -
                                  (1.0 / depth_planes[0]))
  coords_proj = tf.concat([coords_proj, coords_depth_inds], axis=3)
  coords_proj = tf.transpose(
      tf.reshape(coords_proj, [batch_size, cube_res, cube_res, cube_res, 3]),
      [0, 2, 3, 1, 4])
  coords_proj = tf.concat([b[:, :, :, :, tf.newaxis], coords_proj], axis=4)

  # trilinear interpolation gather from MPI
  # interpolate pre-multiplied RGBAs, then un-pre-multiply
  mpi_alpha = mpi[Ellipsis, -1:]
  mpi_channels_p = mpi[Ellipsis, :-1] * mpi_alpha
  mpi_p = tf.concat([mpi_channels_p, mpi_alpha], axis=-1)

  resampled_p = sampling.trilerp_gather(mpi_p, coords_proj, bad_inds)

  resampled_alpha = tf.clip_by_value(resampled_p[Ellipsis, -1:], 0.0, 1.0)
  resampled_channels = resampled_p[Ellipsis, :-1] / (resampled_alpha + 1e-8)
  resampled = tf.concat([resampled_channels, resampled_alpha], axis=-1)

  return resampled, coords_proj
//  [python]
def mpi_resample_cube(mpi, tgt, intrinsics, depth_planes, side_length,
                      cube_res):
  """Resample MPI onto cube centered at target point.

  Args:
    mpi: [B,H,W,D,C], input MPI
    tgt: [B,3], [x,y,z] coordinates for cube center (in reference/mpi frame)
    intrinsics: [B,3,3], MPI reference camera intrinsics
    depth_planes: [D] depth values for MPI planes
    side_length: metric side length of cube
    cube_res: resolution of each cube dimension

  Returns:
    resampled: [B, cube_res, cube_res, cube_res, C]
  """

  batch_size = tf.shape(mpi)[0]
  num_depths = tf.shape(mpi)[3]

  # compute MPI world coordinates
  intrinsics_tile = tf.tile(intrinsics, [num_depths, 1, 1])

  # create cube coordinates
  b_vals = tf.cast(tf.range(batch_size), tf.float32)
  x_vals = tf.linspace(-side_length / 2.0, side_length / 2.0, cube_res)
  y_vals = tf.linspace(-side_length / 2.0, side_length / 2.0, cube_res)
  z_vals = tf.linspace(side_length / 2.0, -side_length / 2.0, cube_res)
  b, y, x, z = tf.meshgrid(b_vals, y_vals, x_vals, z_vals, indexing='ij')

  x = x + tgt[:, 0, tf.newaxis, tf.newaxis, tf.newaxis]
  y = y + tgt[:, 1, tf.newaxis, tf.newaxis, tf.newaxis]
  z = z + tgt[:, 2, tf.newaxis, tf.newaxis, tf.newaxis]

  ones = tf.ones_like(x)
  coords = tf.stack([x, y, z, ones], axis=1)
  coords_r = tf.reshape(
      tf.transpose(coords, [0, 4, 1, 2, 3]),
      [batch_size * cube_res, 4, cube_res, cube_res])

  # store elements with negative z vals for projection
  bad_inds = tf.less(z, 0.0)

  # project into reference camera to transform coordinates into MPI indices
  filler = tf.constant([0.0, 0.0, 0.0, 1.0], shape=[1, 1, 4])
  filler = tf.tile(filler, [batch_size * cube_res, 1, 1])
  intrinsics_tile = tf.tile(intrinsics, [cube_res, 1, 1])
  intrinsics_tile_4 = tf.concat(
      [intrinsics_tile,
       tf.zeros([batch_size * cube_res, 3, 1])], axis=2)
  intrinsics_tile_4 = tf.concat([intrinsics_tile_4, filler], axis=1)
  coords_proj = cam2pixel(coords_r, intrinsics_tile_4)
  coords_depths = tf.transpose(coords_r[:, 2:3, :, :], [0, 2, 3, 1])
  coords_depth_inds = (tf.to_float(num_depths) - 1) * (
      (1.0 / coords_depths) -
      (1.0 / depth_planes[0])) / ((1.0 / depth_planes[-1]) -
                                  (1.0 / depth_planes[0]))
  coords_proj = tf.concat([coords_proj, coords_depth_inds], axis=3)
  coords_proj = tf.transpose(
      tf.reshape(coords_proj, [batch_size, cube_res, cube_res, cube_res, 3]),
      [0, 2, 3, 1, 4])
  coords_proj = tf.concat([b[:, :, :, :, tf.newaxis], coords_proj], axis=4)

  # trilinear interpolation gather from MPI
  # interpolate pre-multiplied RGBAs, then un-pre-multiply
  mpi_alpha = mpi[Ellipsis, -1:]
  mpi_channels_p = mpi[Ellipsis, :-1] * mpi_alpha
  mpi_p = tf.concat([mpi_channels_p, mpi_alpha], axis=-1)

  resampled_p = sampling.trilerp_gather(mpi_p, coords_proj, bad_inds)

  resampled_alpha = tf.clip_by_value(resampled_p[Ellipsis, -1:], 0.0, 1.0)
  resampled_channels = resampled_p[Ellipsis, :-1] / (resampled_alpha + 1e-8)
  resampled = tf.concat([resampled_channels, resampled_alpha], axis=-1)

  return resampled, coords_proj
//  [python]
def _build_losses(self, json_data):
        actor_bound_loss_weight = 10.0
        actor_weight_decay = 0 if (self.ACTOR_WEIGHT_DECAY_KEY not in json_data) else json_data[self.ACTOR_WEIGHT_DECAY_KEY]
        critic_weight_decay = 0 if (self.CRITIC_WEIGHT_DECAY_KEY not in json_data) else json_data[self.CRITIC_WEIGHT_DECAY_KEY]
        
        val_diff = self._tar_val_ph - self._critic_tf
        self._critic_loss_tf = 0.5 * tf.reduce_mean(tf.square(val_diff))

        if (critic_weight_decay != 0):
            self._critic_loss_tf += critic_weight_decay * self._weight_decay_loss(self.MAIN_SCOPE + '/critic')
        
        ratio_tf = tf.exp(self._a_logp_tf - self._old_logp_ph)
        actor_loss0 = self._adv_ph * ratio_tf
        actor_loss1 = self._adv_ph * tf.clip_by_value(ratio_tf, 1.0 - self.ratio_clip, 1 + self.ratio_clip)
        actor_loss_tf = tf.minimum(actor_loss0, actor_loss1)
        self._actor_loss_tf = -tf.reduce_mean(actor_loss_tf)
        
        # for debugging
        self._clip_frac_tf = tf.reduce_mean(tf.to_float(tf.greater(tf.abs(ratio_tf - 1.0), self.ratio_clip)))

        if (actor_bound_loss_weight != 0.0):
            self._actor_loss_tf += actor_bound_loss_weight * self._build_action_bound_loss(self._norm_a_pd_tf)

        if (actor_weight_decay != 0):
            self._actor_loss_tf += actor_weight_decay * self._weight_decay_loss(self.MAIN_SCOPE + '/actor')
         
        return
//  [python]
def _build_losses(self, json_data):
        actor_bound_loss_weight = 10.0
        actor_weight_decay = 0 if (self.ACTOR_WEIGHT_DECAY_KEY not in json_data) else json_data[self.ACTOR_WEIGHT_DECAY_KEY]
        critic_weight_decay = 0 if (self.CRITIC_WEIGHT_DECAY_KEY not in json_data) else json_data[self.CRITIC_WEIGHT_DECAY_KEY]
        
        val_diff = self._tar_val_ph - self._critic_tf
        self._critic_loss_tf = 0.5 * tf.reduce_mean(tf.square(val_diff))

        if (critic_weight_decay != 0):
            self._critic_loss_tf += critic_weight_decay * self._weight_decay_loss(self.MAIN_SCOPE + '/critic')
        
        ratio_tf = tf.exp(self._a_logp_tf - self._old_logp_ph)
        actor_loss0 = self._adv_ph * ratio_tf
        actor_loss1 = self._adv_ph * tf.clip_by_value(ratio_tf, 1.0 - self.ratio_clip, 1 + self.ratio_clip)
        actor_loss_tf = tf.minimum(actor_loss0, actor_loss1)
        self._actor_loss_tf = -tf.reduce_mean(actor_loss_tf)
        
        # for debugging
        self._clip_frac_tf = tf.reduce_mean(tf.cast(tf.greater(tf.abs(ratio_tf - 1.0), self.ratio_clip), dtype=tf.float32))

        if (actor_bound_loss_weight != 0.0):
            self._actor_loss_tf += actor_bound_loss_weight * self._build_action_bound_loss(self._norm_a_pd_tf)

        if (actor_weight_decay != 0):
            self._actor_loss_tf += actor_weight_decay * self._weight_decay_loss(self.MAIN_SCOPE + '/actor')
         
        return
//  [python]
def set_scale_factors_to_mlperf_reference_size(self):
    """Set the parameters to resize the image according to MLPerf reference."""
    # Compute the scale_factor using rounded scaled image size.
    height = tf.shape(self._image)[0]
    width = tf.shape(self._image)[1]
    # Recompute the accurate scale_factor using rounded scaled image size.
    # https://github.com/ddkang/Detectron/blob/80f329530843e66d07ca39e19901d5f3e5daf009/lib/utils/blob.py#L70  # pylint: disable=line-too-long
    min_image_size = tf.to_float(tf.minimum(height, width))
    max_image_size = tf.to_float(tf.maximum(height, width))
    short_side_scale = tf.to_float(self._short_side_image_size) / min_image_size
    long_side_scale = (
        tf.to_float(self._long_side_max_image_size) / max_image_size)
    image_scale = tf.minimum(short_side_scale, long_side_scale)
    scaled_height = tf.to_int32(tf.to_float(height) * image_scale)
    scaled_width = tf.to_int32(tf.to_float(width) * image_scale)
    self._image_scale = image_scale
    self._scaled_height = scaled_height
    self._scaled_width = scaled_width
    return image_scale
//  [python]
def set_scale_factors_to_mlperf_reference_size(self):
    """Set the parameters to resize the image according to MLPerf reference."""
    # Compute the scale_factor using rounded scaled image size.
    height = tf.shape(self._image)[0]
    width = tf.shape(self._image)[1]
    # Recompute the accurate scale_factor using rounded scaled image size.
    # https://github.com/ddkang/Detectron/blob/80f329530843e66d07ca39e19901d5f3e5daf009/lib/utils/blob.py#L70  # pylint: disable=line-too-long
    min_image_size = tf.cast(tf.minimum(height, width), tf.float32)
    max_image_size = tf.cast(tf.maximum(height, width), tf.float32)
    short_side_scale = tf.cast(self._short_side_image_size, tf.float32) / min_image_size
    long_side_scale = (
        tf.cast(self._long_side_max_image_size, tf.float32) / max_image_size)
    image_scale = tf.minimum(short_side_scale, long_side_scale)
    scaled_height = tf.to_int32(tf.float32(height) * image_scale)
    scaled_width = tf.to_int32(tf.float32(width) * image_scale)
    self._image_scale = image_scale
    self._scaled_height = scaled_height
    self._scaled_width = scaled_width
    return image_scale
//  [python]
def _process_metrics(graph, metrics, real_batch_size):
        import tensorflow as tf
        outputs = [real_batch_size]
        val_methods = None
        if metrics is not None:
            idx = 1
            val_methods = []
            for metric_name in metrics:
                metric = metrics[metric_name]
                if tf.is_numeric_tensor(metric):
                    outputs.append(metric)
                    val_methods.append(StatelessMetric(metric_name, idx, 0))
                    idx += 1
                else:
                    outputs += metric.outputs
                    with graph.as_default():
                        val_labels = [tf.identity(v) for v in metric.labels]
                    outputs += val_labels
                    method = TFValidationMethod(metric.val_method,
                                                metric_name,
                                                list(range(idx, idx + len(metric.outputs))),
                                                list(range(idx + len(metric.outputs),
                                                           idx + len(metric.outputs)
                                                           + len(val_labels))))
                    val_methods.append(method)
                    idx += len(metric.outputs) + len(val_labels)

        outputs = [tf.to_float(output) for output in outputs]
        return outputs, val_methods
//  [python]
def _process_metrics(graph, metrics, real_batch_size):
        import tensorflow as tf
        outputs = [real_batch_size]
        val_methods = None
        if metrics is not None:
            idx = 1
            val_methods = []
            for metric_name in metrics:
                metric = metrics[metric_name]
                if tf.is_numeric_tensor(metric):
                    outputs.append(metric)
                    val_methods.append(StatelessMetric(metric_name, idx, 0))
                    idx += 1
                else:
                    outputs += metric.outputs
                    with graph.as_default():
                        val_labels = [tf.identity(v) for v in metric.labels]
                    outputs += val_labels
                    method = TFValidationMethod(metric.val_method,
                                                metric_name,
                                                list(range(idx, idx + len(metric.outputs))),
                                                list(range(idx + len(metric.outputs),
                                                           idx + len(metric.outputs)
                                                           + len(val_labels))))
                    val_methods.append(method)
                    idx += len(metric.outputs) + len(val_labels)

        outputs = [tf.cast(output, tf.float32) for output in outputs]
        return outputs, val_methods
//  [python]
def export(model_dir, loss_tensor, sess, inputs, labels, predictions, grads, variables, graph,
               tensors_with_value, metrics, updates, train_op=None):
        import tensorflow as tf
        with graph.as_default():
            batch_size_tensor = tf.to_float(tf.shape(inputs[0])[0])
        inputs, additional_inputs, additional_values = \
            TFModel._expand_inputs(inputs, tensors_with_value, loss_tensor)
        metric_tensors, val_methods = TFModel._process_metrics(graph, metrics, batch_size_tensor)
        grads = TFModel._process_grads(graph, grads)

        trainable_variables, trainable_variable_placeholders, trainable_assign, \
            extra_variables, extra_variable_assign_placeholders, \
            extra_variable_assign, update_op = \
            TFModel._process_variables(graph, variables, updates)

        meta, saver = \
            TFModel._save_to_dir(model_dir, sess, graph,
                                 metric_tensors,
                                 batch_size_tensor,
                                 loss_tensor, inputs, labels, predictions,
                                 trainable_variables,
                                 trainable_variable_placeholders,
                                 trainable_assign,
                                 extra_variables,
                                 extra_variable_assign_placeholders,
                                 extra_variable_assign,
                                 grads, update_op, train_op,
                                 additional_inputs,
                                 additional_values)
        return meta, saver, val_methods
//  [python]
def export(model_dir, loss_tensor, sess, inputs, labels, predictions, grads, variables, graph,
               tensors_with_value, metrics, updates, train_op=None):
        import tensorflow as tf
        with graph.as_default():
            batch_size_tensor = tf.cast(tf.shape(inputs[0])[0], tf.float32)
        inputs, additional_inputs, additional_values = \
            TFModel._expand_inputs(inputs, tensors_with_value, loss_tensor)
        metric_tensors, val_methods = TFModel._process_metrics(graph, metrics, batch_size_tensor)
        grads = TFModel._process_grads(graph, grads)

        trainable_variables, trainable_variable_placeholders, trainable_assign, \
            extra_variables, extra_variable_assign_placeholders, \
            extra_variable_assign, update_op = \
            TFModel._process_variables(graph, variables, updates)

        meta, saver = \
            TFModel._save_to_dir(model_dir, sess, graph,
                                 metric_tensors,
                                 batch_size_tensor,
                                 loss_tensor, inputs, labels, predictions,
                                 trainable_variables,
                                 trainable_variable_placeholders,
                                 trainable_assign,
                                 extra_variables,
                                 extra_variable_assign_placeholders,
                                 extra_variable_assign,
                                 grads, update_op, train_op,
                                 additional_inputs,
                                 additional_values)
        return meta, saver, val_methods
//  [python]
def _process_variables(graph, variables, updates):
        import tensorflow as tf
        all_trainable_variables = variables

        name2idx = dict([(v.name, idx) for idx, v in enumerate(all_trainable_variables)])

        all_variables = graph.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)

        update_ops = graph.get_collection(tf.GraphKeys.UPDATE_OPS)

        if updates is not None:
            update_ops += updates

        trainable_variables = [0] * len(all_trainable_variables)
        trainable_assigns = [0] * len(all_trainable_variables)
        trainable_variable_placeholders = [0] * len(all_trainable_variables)
        extra_variables = []
        extra_variable_assigns = []
        extra_variable_assign_placeholders = []
        for v in all_variables:
            p = tf.placeholder(dtype=v.dtype, shape=v.shape)
            a = tf.assign(v, p)

            # special treatment for ResourceVariable
            if v.op.type == "VarHandleOp":
                v_float_value = tf.to_float(v.read_value())
            else:
                v_float_value = tf.to_float(v)

            if v.name in name2idx:
                trainable_variables[name2idx[v.name]] = v_float_value
                trainable_assigns[name2idx[v.name]] = a
                trainable_variable_placeholders[name2idx[v.name]] = p
            else:
                extra_variables.append(v_float_value)
                extra_variable_assigns.append(a)
                extra_variable_assign_placeholders.append(p)

        extra_variable_assign = tf.group(*extra_variable_assigns)
        trainable_assign = tf.group(*trainable_assigns)
        update_op = tf.group(update_ops)

        return trainable_variables, trainable_variable_placeholders, trainable_assign, \
            extra_variables, extra_variable_assign_placeholders, \
            extra_variable_assign, update_op
//  [python]
def _process_variables(graph, variables, updates):
        import tensorflow as tf
        all_trainable_variables = variables

        name2idx = dict([(v.name, idx) for idx, v in enumerate(all_trainable_variables)])

        all_variables = graph.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)

        update_ops = graph.get_collection(tf.GraphKeys.UPDATE_OPS)

        if updates is not None:
            update_ops += updates

        trainable_variables = [0] * len(all_trainable_variables)
        trainable_assigns = [0] * len(all_trainable_variables)
        trainable_variable_placeholders = [0] * len(all_trainable_variables)
        extra_variables = []
        extra_variable_assigns = []
        extra_variable_assign_placeholders = []
        for v in all_variables:
            p = tf.placeholder(dtype=v.dtype, shape=v.shape)
            a = tf.assign(v, p)

            # special treatment for ResourceVariable
            if v.op.type == "VarHandleOp":
                v_float_value = tf.cast(v.read_value(), tf.float32)
            else:
                v_float_value = tf.cast(v, tf.float32)

            if v.name in name2idx:
                trainable_variables[name2idx[v.name]] = v_float_value
                trainable_assigns[name2idx[v.name]] = a
                trainable_variable_placeholders[name2idx[v.name]] = p
            else:
                extra_variables.append(v_float_value)
                extra_variable_assigns.append(a)
                extra_variable_assign_placeholders.append(p)

        extra_variable_assign = tf.group(*extra_variable_assigns)
        trainable_assign = tf.group(*trainable_assigns)
        update_op = tf.group(update_ops)

        return trainable_variables, trainable_variable_placeholders, trainable_assign, \
            extra_variables, extra_variable_assign_placeholders, \
            extra_variable_assign, update_op
//  [python]
def _compute_f1(hypes, labels, softmax, epsilon):
    labels = tf.to_float(tf.reshape(labels, (-1, 2)))[:, 1]
    logits = softmax[:, 1]
    true_positive = tf.reduce_sum(labels*logits)
    false_positive = tf.reduce_sum((1-labels)*logits)

    recall = true_positive / tf.reduce_sum(labels)
    precision = true_positive / (true_positive + false_positive + epsilon)

    score = 2*recall * precision / (precision + recall)
    f1_score = 1 - 2*recall * precision / (precision + recall)

    return f1_score
//  [python]
def _compute_f1(hypes, labels, softmax, epsilon):
    labels = tf.cast(tf.reshape(labels, (-1, 2)), tf.float32)[:, 1]
    logits = softmax[:, 1]
    true_positive = tf.reduce_sum(labels*logits)
    false_positive = tf.reduce_sum((1-labels)*logits)

    recall = true_positive / tf.reduce_sum(labels)
    precision = true_positive / (true_positive + false_positive + epsilon)

    score = 2*recall * precision / (precision + recall)
    f1_score = 1 - 2*recall * precision / (precision + recall)

    return f1_score
//  [python]
def loss(hypes, decoded_logits, labels):
    """Calculate the loss from the logits and the labels.

    Args:
      logits: Logits tensor, float - [batch_size, NUM_CLASSES].
      labels: Labels tensor, int32 - [batch_size].

    Returns:
      loss: Loss tensor of type float.
    """
    logits = decoded_logits['logits']
    with tf.name_scope('loss'):

        logits = tf.reshape(logits, (-1, 2))
        shape = [logits.get_shape()[0], 2]
        epsilon = tf.constant(value=hypes['solver']['epsilon'])
        # logits = logits + epsilon
        labels = tf.to_float(tf.reshape(labels, (-1, 2)))

        softmax = tf.nn.softmax(logits) + epsilon

        if hypes['loss'] == 'xentropy':
            cross_entropy_mean = _compute_cross_entropy_mean(hypes, labels,
                                                             softmax)
        elif hypes['loss'] == 'softF1':
            cross_entropy_mean = _compute_f1(hypes, labels, softmax, epsilon)

        elif hypes['loss'] == 'softIU':
            cross_entropy_mean = _compute_soft_ui(hypes, labels, softmax,
                                                  epsilon)

        reg_loss_col = tf.GraphKeys.REGULARIZATION_LOSSES

        weight_loss = tf.add_n(tf.get_collection(reg_loss_col),
                               name='reg_loss')

        total_loss = cross_entropy_mean + weight_loss

        losses = {}
        losses['total_loss'] = total_loss
        losses['xentropy'] = cross_entropy_mean
        losses['weight_loss'] = weight_loss

    return losses
//  [python]
def loss(hypes, decoded_logits, labels):
    """Calculate the loss from the logits and the labels.

    Args:
      logits: Logits tensor, float - [batch_size, NUM_CLASSES].
      labels: Labels tensor, int32 - [batch_size].

    Returns:
      loss: Loss tensor of type float.
    """
    logits = decoded_logits['logits']
    with tf.name_scope('loss'):

        logits = tf.reshape(logits, (-1, 2))
        shape = [logits.get_shape()[0], 2]
        epsilon = tf.constant(value=hypes['solver']['epsilon'])
        # logits = logits + epsilon
        labels = tf.cast(tf.reshape(labels, (-1, 2)), tf.float32)

        softmax = tf.nn.softmax(logits) + epsilon

        if hypes['loss'] == 'xentropy':
            cross_entropy_mean = _compute_cross_entropy_mean(hypes, labels,
                                                             softmax)
        elif hypes['loss'] == 'softF1':
            cross_entropy_mean = _compute_f1(hypes, labels, softmax, epsilon)

        elif hypes['loss'] == 'softIU':
            cross_entropy_mean = _compute_soft_ui(hypes, labels, softmax,
                                                  epsilon)

        reg_loss_col = tf.GraphKeys.REGULARIZATION_LOSSES

        weight_loss = tf.add_n(tf.get_collection(reg_loss_col),
                               name='reg_loss')

        total_loss = cross_entropy_mean + weight_loss

        losses = {}
        losses['total_loss'] = total_loss
        losses['xentropy'] = cross_entropy_mean
        losses['weight_loss'] = weight_loss

    return losses
//  [python]
def input_fn_builder(input_file, seq_length, is_training,
                     drop_remainder, use_tpu, bsz, is_v2):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""

  name_to_features = {
      "unique_ids": tf.FixedLenFeature([], tf.int64),
      "input_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "input_mask": tf.FixedLenFeature([seq_length], tf.int64),
      "segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
  }
  # p_mask is not required for SQuAD v1.1
  if is_v2:
    name_to_features["p_mask"] = tf.FixedLenFeature([seq_length], tf.int64)

  if is_training:
    name_to_features["start_positions"] = tf.FixedLenFeature([], tf.int64)
    name_to_features["end_positions"] = tf.FixedLenFeature([], tf.int64)
    name_to_features["is_impossible"] = tf.FixedLenFeature([], tf.int64)

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

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.to_int32(t)
      example[name] = t

    return example

  def input_fn(params):
    """The actual input function."""
    if use_tpu:
      batch_size = params["batch_size"]
    else:
      batch_size = bsz

    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    d = tf.data.TFRecordDataset(input_file)
    if is_training:
      d = d.repeat()
      d = d.shuffle(buffer_size=100)

    d = d.apply(
        tf.data.experimental.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            drop_remainder=drop_remainder))

    return d

  return input_fn
//  [python]
def input_fn_builder(input_file, seq_length, is_training,
                     drop_remainder, use_tpu, bsz, is_v2):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""

  name_to_features = {
      "unique_ids": tf.FixedLenFeature([], tf.int64),
      "input_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "input_mask": tf.FixedLenFeature([seq_length], tf.int64),
      "segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
  }
  # p_mask is not required for SQuAD v1.1
  if is_v2:
    name_to_features["p_mask"] = tf.FixedLenFeature([seq_length], tf.int64)

  if is_training:
    name_to_features["start_positions"] = tf.FixedLenFeature([], tf.int64)
    name_to_features["end_positions"] = tf.FixedLenFeature([], tf.int64)
    name_to_features["is_impossible"] = tf.FixedLenFeature([], tf.int64)

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

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.cast(t, tf.int32)
      example[name] = t

    return example

  def input_fn(params):
    """The actual input function."""
    if use_tpu:
      batch_size = params["batch_size"]
    else:
      batch_size = bsz

    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    d = tf.data.TFRecordDataset(input_file)
    if is_training:
      d = d.repeat()
      d = d.shuffle(buffer_size=100)

    d = d.apply(
        tf.data.experimental.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            drop_remainder=drop_remainder))

    return d

  return input_fn
//  [python]
def evaluation(hyp, images, labels, decoded_logits, losses, global_step):
    """Evaluate the quality of the logits at predicting the label.

    Args:
      logits: Logits tensor, float - [batch_size, NUM_CLASSES].
      labels: Labels tensor, int32 - [batch_size], with values in the
        range [0, NUM_CLASSES).

    Returns:
      A scalar int32 tensor with the number of examples (out of batch_size)
      that were predicted correctly.
    """
    # For a classifier model, we can use the in_top_k Op.
    # It returns a bool tensor with shape [batch_size] that is true for
    # the examples where the label's is was in the top k (here k=1)
    # of all logits for that example.
    eval_list = []
    logits = tf.reshape(decoded_logits['logits'], (-1, 2))
    labels = tf.reshape(labels, (-1, 2))

    pred = tf.argmax(logits, dimension=1)

    negativ = tf.to_int32(tf.equal(pred, 0))
    tn = tf.reduce_sum(negativ*labels[:, 0])
    fn = tf.reduce_sum(negativ*labels[:, 1])

    positive = tf.to_int32(tf.equal(pred, 1))
    tp = tf.reduce_sum(positive*labels[:, 1])
    fp = tf.reduce_sum(positive*labels[:, 0])

    eval_list.append(('Acc. ', (tn+tp)/(tn + fn + tp + fp)))
    eval_list.append(('xentropy', losses['xentropy']))
    eval_list.append(('weight_loss', losses['weight_loss']))

    # eval_list.append(('Precision', tp/(tp + fp)))
    # eval_list.append(('True BG', tn/(tn + fp)))
    # eval_list.append(('True Street [Recall]', tp/(tp + fn)))

    return eval_list
//  [python]
def evaluation(hyp, images, labels, decoded_logits, losses, global_step):
    """Evaluate the quality of the logits at predicting the label.

    Args:
      logits: Logits tensor, float - [batch_size, NUM_CLASSES].
      labels: Labels tensor, int32 - [batch_size], with values in the
        range [0, NUM_CLASSES).

    Returns:
      A scalar int32 tensor with the number of examples (out of batch_size)
      that were predicted correctly.
    """
    # For a classifier model, we can use the in_top_k Op.
    # It returns a bool tensor with shape [batch_size] that is true for
    # the examples where the label's is was in the top k (here k=1)
    # of all logits for that example.
    eval_list = []
    logits = tf.reshape(decoded_logits['logits'], (-1, 2))
    labels = tf.reshape(labels, (-1, 2))

    pred = tf.argmax(logits, dimension=1)

    negativ = tf.cast(tf.equal(pred, 0), tf.int32)
    tn = tf.reduce_sum(negativ*labels[:, 0])
    fn = tf.reduce_sum(negativ*labels[:, 1])

    positive = tf.cast(tf.equal(pred, 1), tf.int32)
    tp = tf.reduce_sum(positive*labels[:, 1])
    fp = tf.reduce_sum(positive*labels[:, 0])

    eval_list.append(('Acc. ', (tn+tp)/(tn + fn + tp + fp)))
    eval_list.append(('xentropy', losses['xentropy']))
    eval_list.append(('weight_loss', losses['weight_loss']))

    # eval_list.append(('Precision', tp/(tp + fp)))
    # eval_list.append(('True BG', tn/(tn + fp)))
    # eval_list.append(('True Street [Recall]', tp/(tp + fn)))

    return eval_list
//  [python]
def file_based_input_fn_builder(input_file, seq_length, is_training,
                                drop_remainder, num_labels):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""

  name_to_features = {
      "input_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "input_mask": tf.FixedLenFeature([seq_length], tf.int64),
      "segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "label_ids": tf.FixedLenFeature([num_labels], tf.float32),
      "is_real_example": tf.FixedLenFeature([], tf.int64),
  }

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

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.to_int32(t)
      example[name] = t

    return example

  def input_fn(params):
    """The actual input function."""
    batch_size = params["batch_size"]

    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    d = tf.data.TFRecordDataset(input_file)
    if is_training:
      d = d.repeat()
      d = d.shuffle(buffer_size=100)

    d = d.apply(
        tf.data.experimental.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            drop_remainder=drop_remainder))

    return d

  return input_fn
//  [python]
def file_based_input_fn_builder(input_file, seq_length, is_training,
                                drop_remainder, num_labels):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""

  name_to_features = {
      "input_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "input_mask": tf.FixedLenFeature([seq_length], tf.int64),
      "segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
      "label_ids": tf.FixedLenFeature([num_labels], tf.float32),
      "is_real_example": tf.FixedLenFeature([], tf.int64),
  }

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

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.cast(t, tf.int32)
      example[name] = t

    return example

  def input_fn(params):
    """The actual input function."""
    batch_size = params["batch_size"]

    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    d = tf.data.TFRecordDataset(input_file)
    if is_training:
      d = d.repeat()
      d = d.shuffle(buffer_size=100)

    d = d.apply(
        tf.data.experimental.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            drop_remainder=drop_remainder))

    return d

  return input_fn
//  [python]
def get_dropout(self, dropout_rate, is_training):
        return 1 - (tf.to_float(is_training) * dropout_rate)
//  [python]
def get_dropout(self, dropout_rate, is_training):
        return 1 - (tf.cast(is_training, tf.float32) * dropout_rate)
//  [python]
def file_based_input_fn_builder(input_file, seq_length, is_training, drop_remainder):
    name_to_features = {
        "input_ids": tf.FixedLenFeature([seq_length], tf.int64),
        "input_mask": tf.FixedLenFeature([seq_length], tf.int64),
        "segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
        "label_ids": tf.FixedLenFeature([seq_length], tf.int64),
        # "label_ids":tf.VarLenFeature(tf.int64),
        # "label_mask": tf.FixedLenFeature([seq_length], tf.int64),
    }

    def _decode_record(record, name_to_features):
        example = tf.parse_single_example(record, name_to_features)
        for name in list(example.keys()):
            t = example[name]
            if t.dtype == tf.int64:
                t = tf.to_int32(t)
            example[name] = t
        return example

    def input_fn(params):
        batch_size = params["batch_size"]
        d = tf.data.TFRecordDataset(input_file)
        if is_training:
            d = d.repeat()
            d = d.shuffle(buffer_size=100)
        d = d.apply(tf.contrib.data.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            drop_remainder=drop_remainder
        ))
        return d

    return input_fn
//  [python]
def file_based_input_fn_builder(input_file, seq_length, is_training, drop_remainder):
    name_to_features = {
        "input_ids": tf.FixedLenFeature([seq_length], tf.int64),
        "input_mask": tf.FixedLenFeature([seq_length], tf.int64),
        "segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
        "label_ids": tf.FixedLenFeature([seq_length], tf.int64),
        # "label_ids":tf.VarLenFeature(tf.int64),
        # "label_mask": tf.FixedLenFeature([seq_length], tf.int64),
    }

    def _decode_record(record, name_to_features):
        example = tf.parse_single_example(record, name_to_features)
        for name in list(example.keys()):
            t = example[name]
            if t.dtype == tf.int64:
                t = tf.cast(t, tf.int32)
            example[name] = t
        return example

    def input_fn(params):
        batch_size = params["batch_size"]
        d = tf.data.TFRecordDataset(input_file)
        if is_training:
            d = d.repeat()
            d = d.shuffle(buffer_size=100)
        d = d.apply(tf.contrib.data.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            drop_remainder=drop_remainder
        ))
        return d

    return input_fn
//  [python]
def get_candidate_labels(self, candidate_starts, candidate_ends, labeled_starts, labeled_ends, labels):
        same_start = tf.equal(tf.expand_dims(labeled_starts, 1), tf.expand_dims(candidate_starts, 0)) # [num_labeled, num_candidates]
        same_end = tf.equal(tf.expand_dims(labeled_ends, 1), tf.expand_dims(candidate_ends, 0)) # [num_labeled, num_candidates]
        same_span = tf.logical_and(same_start, same_end) # [num_labeled, num_candidates]
        candidate_labels = tf.matmul(tf.expand_dims(labels, 0), tf.to_int32(same_span)) # [1, num_candidates]
        candidate_labels = tf.squeeze(candidate_labels, 0) # [num_candidates]
        return candidate_labels
//  [python]
def get_candidate_labels(self, candidate_starts, candidate_ends, labeled_starts, labeled_ends, labels):
        same_start = tf.equal(tf.expand_dims(labeled_starts, 1), tf.expand_dims(candidate_starts, 0)) # [num_labeled, num_candidates]
        same_end = tf.equal(tf.expand_dims(labeled_ends, 1), tf.expand_dims(candidate_ends, 0)) # [num_labeled, num_candidates]
        same_span = tf.logical_and(same_start, same_end) # [num_labeled, num_candidates]
        candidate_labels = tf.matmul(tf.expand_dims(labels, 0), tf.cast(same_span, tf.int32)) # [1, num_candidates]
        candidate_labels = tf.squeeze(candidate_labels, 0) # [num_candidates]
        return candidate_labels
//  [python]
def predict_depth(self, rgb, sensor_depth=None):
    del sensor_depth  # unused
    with tf.variable_scope('depth_prediction', reuse=tf.AUTO_REUSE):
      if self._mode == tf.estimator.ModeKeys.TRAIN:
        noise_stddev = 0.5
        global_step = tf.train.get_global_step()
        rampup_steps = self._params['layer_norm_noise_rampup_steps']
        if global_step is not None and rampup_steps > 0:
          # If global_step is available, ramp up the noise.
          noise_stddev *= tf.square(
              tf.minimum(tf.to_float(global_step) / float(rampup_steps), 1.0))
      else:
        noise_stddev = 0.0

      def _normalizer_fn(x, is_train, name='bn'):
        return randomized_layer_norm(
            x, is_train=is_train, name=name, stddev=noise_stddev)

      if self._params['learn_scale']:
        depth_scale = tf.get_variable('depth_scale', initializer=1.0)
        maybe_summary.scalar('depth_scale', depth_scale)
      else:
        depth_scale = 1.0

      return depth_scale * depth_prediction_resnet18unet(
          2 * rgb - 1.0,
          self._mode == tf.estimator.ModeKeys.TRAIN,
          self._params['weight_decay'],
          _normalizer_fn,
          reflect_padding=self._params['reflect_padding'])
//  [python]
def predict_depth(self, rgb, sensor_depth=None):
    del sensor_depth  # unused
    with tf.variable_scope('depth_prediction', reuse=tf.AUTO_REUSE):
      if self._mode == tf.estimator.ModeKeys.TRAIN:
        noise_stddev = 0.5
        global_step = tf.train.get_global_step()
        rampup_steps = self._params['layer_norm_noise_rampup_steps']
        if global_step is not None and rampup_steps > 0:
          # If global_step is available, ramp up the noise.
          noise_stddev *= tf.square(
              tf.minimum(tf.cast(global_step, tf.float32) / float(rampup_steps), 1.0))
      else:
        noise_stddev = 0.0

      def _normalizer_fn(x, is_train, name='bn'):
        return randomized_layer_norm(
            x, is_train=is_train, name=name, stddev=noise_stddev)

      if self._params['learn_scale']:
        depth_scale = tf.get_variable('depth_scale', initializer=1.0)
        maybe_summary.scalar('depth_scale', depth_scale)
      else:
        depth_scale = 1.0

      return depth_scale * depth_prediction_resnet18unet(
          2 * rgb - 1.0,
          self._mode == tf.estimator.ModeKeys.TRAIN,
          self._params['weight_decay'],
          _normalizer_fn,
          reflect_padding=self._params['reflect_padding'])