jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def van(first_enc,
        first_frame,
        current_enc,
        gt_image,
        reuse=False,
        scope_prefix='',
        hparams=None):
  """Implements a VAN.

  Args:
    first_enc: The first encoding.
    first_frame: The first ground truth frame.
    current_enc: The encoding of the frame to generate.
    gt_image: The ground truth image, only used for regularization.
    reuse: To reuse in variable scope or not.
    scope_prefix: The prefix before the scope name.
    hparams: The python hparams.

  Returns:
    The generated image.
  """
  with tf.variable_scope(scope_prefix + 'van', reuse=reuse):
    output_shape = first_frame.get_shape().as_list()
    output_shape[0] = -1

    first_depth = 64

    f_first_enc, _ = van_enc_2d(first_enc, first_depth)
    f_first_frame, image_enc_history = van_image_enc_2d(
        first_frame, first_depth, hparams=hparams)
    f_current_enc, van_higher_level = van_enc_2d(
        current_enc, first_depth, reuse=True)
    f_gt_image, _ = van_image_enc_2d(gt_image, first_depth, True,
                                     hparams=hparams)

    analogy_t = analogy_computation_2d(
        f_first_enc, f_first_frame, f_current_enc, first_depth)
    enc_img = f_current_enc + analogy_t

    img = van_dec_2d(
        enc_img, image_enc_history, output_shape, first_depth, hparams=hparams)

    batch_size = tf.to_float(tf.shape(first_enc)[0])
    r_loss = tf.nn.l2_loss(f_gt_image - f_current_enc - analogy_t) / batch_size

    return img, r_loss, van_higher_level
//  [python]
def van(first_enc,
        first_frame,
        current_enc,
        gt_image,
        reuse=False,
        scope_prefix='',
        hparams=None):
  """Implements a VAN.

  Args:
    first_enc: The first encoding.
    first_frame: The first ground truth frame.
    current_enc: The encoding of the frame to generate.
    gt_image: The ground truth image, only used for regularization.
    reuse: To reuse in variable scope or not.
    scope_prefix: The prefix before the scope name.
    hparams: The python hparams.

  Returns:
    The generated image.
  """
  with tf.variable_scope(scope_prefix + 'van', reuse=reuse):
    output_shape = first_frame.get_shape().as_list()
    output_shape[0] = -1

    first_depth = 64

    f_first_enc, _ = van_enc_2d(first_enc, first_depth)
    f_first_frame, image_enc_history = van_image_enc_2d(
        first_frame, first_depth, hparams=hparams)
    f_current_enc, van_higher_level = van_enc_2d(
        current_enc, first_depth, reuse=True)
    f_gt_image, _ = van_image_enc_2d(gt_image, first_depth, True,
                                     hparams=hparams)

    analogy_t = analogy_computation_2d(
        f_first_enc, f_first_frame, f_current_enc, first_depth)
    enc_img = f_current_enc + analogy_t

    img = van_dec_2d(
        enc_img, image_enc_history, output_shape, first_depth, hparams=hparams)

    batch_size = tf.cast(tf.shape(first_enc)[0], tf.float32)
    r_loss = tf.nn.l2_loss(f_gt_image - f_current_enc - analogy_t) / batch_size

    return img, r_loss, van_higher_level
//  [python]
def calc_loss_psnr(gen_images, images, name, hparams=None, use_l1_loss=False):
  """Calculates loss and psnr for predictions over multiple timesteps."""
  del hparams
  with tf.name_scope(name):
    loss, error, psnr_all = 0.0, 0.0, 0.0
    for _, x, gx in zip(range(len(gen_images)), images, gen_images):
      recon_cost = mean_squared_error(x, gx)
      if use_l1_loss:
        recon_cost = l1_error(x, gx)

      error_i = l1_error(x, gx)
      psnr_i = peak_signal_to_noise_ratio(x, gx)
      psnr_all += psnr_i
      error += error_i
      loss += recon_cost

    psnr_all /= tf.to_float(len(gen_images))
    loss /= tf.to_float(len(gen_images))
    error /= tf.to_float(len(gen_images))

    # if not hparams.use_tpu:
    tf.summary.scalar('psnr_all', psnr_all)
    tf.summary.scalar('loss', loss)

    return loss, psnr_all
//  [python]
def calc_loss_psnr(gen_images, images, name, hparams=None, use_l1_loss=False):
  """Calculates loss and psnr for predictions over multiple timesteps."""
  del hparams
  with tf.name_scope(name):
    loss, error, psnr_all = 0.0, 0.0, 0.0
    for _, x, gx in zip(range(len(gen_images)), images, gen_images):
      recon_cost = mean_squared_error(x, gx)
      if use_l1_loss:
        recon_cost = l1_error(x, gx)

      error_i = l1_error(x, gx)
      psnr_i = peak_signal_to_noise_ratio(x, gx)
      psnr_all += psnr_i
      error += error_i
      loss += recon_cost

    psnr_all /= tf.cast(len(gen_images), tf.float32)
    loss /= tf.cast(len(gen_images), tf.float32)
    error /= tf.cast(len(gen_images), tf.float32)

    # if not hparams.use_tpu:
    tf.summary.scalar('psnr_all', psnr_all)
    tf.summary.scalar('loss', loss)

    return loss, psnr_all
//  [python]
def _scale_preserving_resize(image):
    target_size = tf.convert_to_tensor(config['preprocessing']['resize'])
    scales = tf.to_float(tf.divide(target_size, tf.shape(image)[:2]))
    new_size = tf.to_float(tf.shape(image)[:2]) * tf.reduce_max(scales)
    image = tf.image.resize_images(image, tf.to_int32(new_size),
                                   method=tf.image.ResizeMethod.BILINEAR)
    return tf.image.resize_image_with_crop_or_pad(image, target_size[0],
                                                  target_size[1])
//  [python]
def _scale_preserving_resize(image):
    target_size = tf.convert_to_tensor(config['preprocessing']['resize'])
    scales = tf.cast(tf.divide(target_size, tf.shape(image)[:2]), tf.float32)
    new_size = tf.cast(tf.shape(image)[:2], tf.float32) * tf.reduce_max(scales)
    image = tf.image.resize_images(image, tf.cast(new_size, tf.int32),
                                   method=tf.image.ResizeMethod.BILINEAR)
    return tf.image.resize_image_with_crop_or_pad(image, target_size[0],
                                                  target_size[1])
//  [python]
def crop_or_pad(waves, length, channels):
  """Crop or pad wave to have shape [N, length, channels].
  Args:
    waves: A 3D `Tensor` of NLC format.
    length: A Python scalar. The output wave size.
    channels: Number of output waves channels.
  Returns:
    A 3D `Tensor` of NLC format with shape [N, length, channels].
  """
  waves = tf.convert_to_tensor(waves)
  batch_size = int(waves.shape[0])
  waves_shape = tf.shape(waves)

  # Force audio length.
  pad = tf.maximum(0, length - waves_shape[1])
  right_pad = tf.to_int32(tf.to_float(pad) / 2.0)
  left_pad = pad - right_pad
  waves = tf.pad(waves, [[0, 0], [left_pad, right_pad], [0, 0]])
  waves = waves[:, :length, :]

  # Force number of channels.
  num_repeats = tf.to_int32(
      tf.ceil(tf.to_float(channels) / tf.to_float(waves_shape[2])))
  waves = tf.tile(waves, [1, 1, num_repeats])[:, :, :channels]

  waves.set_shape([batch_size, length, channels])
  return waves
//  [python]
def crop_or_pad(waves, length, channels):
  """Crop or pad wave to have shape [N, length, channels].
  Args:
    waves: A 3D `Tensor` of NLC format.
    length: A Python scalar. The output wave size.
    channels: Number of output waves channels.
  Returns:
    A 3D `Tensor` of NLC format with shape [N, length, channels].
  """
  waves = tf.convert_to_tensor(waves)
  batch_size = int(waves.shape[0])
  waves_shape = tf.shape(waves)

  # Force audio length.
  pad = tf.maximum(0, length - waves_shape[1])
  right_pad = tf.cast(tf.math.ceil(tf.cast(pad, tf.float32) / 2.0), tf.int32)
  left_pad = pad - right_pad
  waves = tf.pad(waves, [[0, 0], [left_pad, right_pad], [0, 0]])
  waves = waves[:, :length, :]

  # Force number of channels.
  num_repeats = tf.cast(
      tf.math.ceil(tf.cast(channels, tf.float32) / tf.cast(waves_shape[2], tf.float32)), tf.int32)
  waves = tf.tile(waves, [1, 1, num_repeats])[:, :, :channels]

  waves.set_shape([batch_size, length, channels])
  return waves
//  [python]
def preprocess_image(image,
                     output_height,
                     output_width,
                     is_training,
                     use_grayscale=False):
  """Preprocesses the given image.

  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.
    is_training: `True` if we're preprocessing the image for training and
      `False` otherwise.
    use_grayscale: Whether to convert the image from RGB to grayscale.

  Returns:
    A preprocessed image.
  """
  del is_training  # Unused argument
  image = tf.to_float(image)
  if use_grayscale:
    image = tf.image.rgb_to_grayscale(image)
  image = tf.image.resize_image_with_crop_or_pad(
      image, output_width, output_height)
  image = tf.subtract(image, 128.0)
  image = tf.div(image, 128.0)
  return image
//  [python]
def preprocess_image(image,
                     output_height,
                     output_width,
                     is_training,
                     use_grayscale=False):
  """Preprocesses the given image.

  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.
    is_training: `True` if we're preprocessing the image for training and
      `False` otherwise.
    use_grayscale: Whether to convert the image from RGB to grayscale.

  Returns:
    A preprocessed image.
  """
  del is_training  # Unused argument
  image = tf.cast(image, tf.float32)
  if use_grayscale:
    image = tf.image.rgb_to_grayscale(image)
  image = tf.image.resize_image_with_crop_or_pad(
      image, output_width, output_height)
  image = tf.subtract(image, 128.0)
  image = tf.div(image, 128.0)
  return image
//  [python]
def __init__(self):
        self.n_visible = 17770
        self.batch_size = 128
        self.n_hidden = 100
        self.momentum = tf.constant(0.9)
        self.weight_decay = tf.constant(0.0)
        self.k = 1
        self.num_rat = 5
        self.lr_weights = tf.constant(0.0015)
        self.lr_vb = tf.constant(0.0015)
        self.lr_hb = tf.constant(0.1)

        self.anneal = False
        self.anneal_val = tf.constant(0.0)
        self.num_users = 458293

        # if momentum:
        #     self.momentum = tf.placeholder(tf.float32)
        # else:
        #     self.momentum = 0.0
        # Initialize weights and biases
        initial_weights = tf.random.normal([self.n_visible, self.num_rat, self.n_hidden], stddev=0.01)
        # initial_weights = tf.reduce_mean(initial_weights, axis=1)
        # initial_weights = tf.stack([initial_weights, initial_weights, initial_weights, initial_weights, initial_weights], axis=1)
        self.weights = tfe.Variable(initial_weights, name="weights")


        arr = np.loadtxt("movie_frequencies.dta")

        visible_bias = tf.to_float(tf.transpose(tf.constant(arr)))
        self.hidden_bias = tfe.Variable(tf.constant(0.0, shape=[self.num_users, self.n_hidden]), name='h_bias')
        self.visible_bias = tfe.Variable(visible_bias, name='v_bias')

        self.weight_v = tf.zeros(tf.shape(self.weights))
        self.visible_bias_v = tf.zeros(tf.shape(self.visible_bias))
        self.hidden_bias_v = tf.zeros(tf.shape(self.hidden_bias))
//  [python]
def __init__(self):
        self.n_visible = 17770
        self.batch_size = 128
        self.n_hidden = 100
        self.momentum = tf.constant(0.9)
        self.weight_decay = tf.constant(0.0)
        self.k = 1
        self.num_rat = 5
        self.lr_weights = tf.constant(0.0015)
        self.lr_vb = tf.constant(0.0015)
        self.lr_hb = tf.constant(0.1)

        self.anneal = False
        self.anneal_val = tf.constant(0.0)
        self.num_users = 458293

        # if momentum:
        #     self.momentum = tf.placeholder(tf.float32)
        # else:
        #     self.momentum = 0.0
        # Initialize weights and biases
        initial_weights = tf.random.normal([self.n_visible, self.num_rat, self.n_hidden], stddev=0.01)
        # initial_weights = tf.reduce_mean(initial_weights, axis=1)
        # initial_weights = tf.stack([initial_weights, initial_weights, initial_weights, initial_weights, initial_weights], axis=1)
        self.weights = tfe.Variable(initial_weights, name="weights")


        arr = np.loadtxt("movie_frequencies.dta")

        visible_bias = tf.cast(tf.transpose(tf.constant(arr)), tf.float32)
        self.hidden_bias = tfe.Variable(tf.constant(0.0, shape=[self.num_users, self.n_hidden]), name='h_bias')
        self.visible_bias = tfe.Variable(visible_bias, name='v_bias')

        self.weight_v = tf.zeros(tf.shape(self.weights))
        self.visible_bias_v = tf.zeros(tf.shape(self.visible_bias))
        self.hidden_bias_v = tf.zeros(tf.shape(self.hidden_bias))
//  [python]
def make_mnist(config):
  """Creates the MNIST dataset."""

  def to_float(x):
    return tf.to_float(x) / 255.

  transform = [to_float]

  if config.canvas_size != 28:
    transform.append(functools.partial(preprocess.pad_and_shift,
                                       output_size=config.canvas_size,
                                       shift=None))

  batch_size = config.batch_size
  res = AttrDict(
      trainset=image.create(
          'mnist', subset='train', batch_size=batch_size, transforms=transform),
      validset=image.create(
          'mnist', subset='test', batch_size=batch_size, transforms=transform))

  return res
//  [python]
def make_mnist(config):
  """Creates the MNIST dataset."""

  def to_float(x):
    return tf.cast(x, tf.float32) / 255.

  transform = [to_float]

  if config.canvas_size != 28:
    transform.append(functools.partial(preprocess.pad_and_shift,
                                       output_size=config.canvas_size,
                                       shift=None))

  batch_size = config.batch_size
  res = AttrDict(
      trainset=image.create(
          'mnist', subset='train', batch_size=batch_size, transforms=transform),
      validset=image.create(
          'mnist', subset='test', batch_size=batch_size, transforms=transform))

  return res
//  [python]
def _apply_negative_infinity_mask(tensor, mask):
  """Where mask is true, add a large negative value."""
  tensor += tf.to_float(mask) * -INF
  tensor = tf.maximum(tensor, -INF)
  return tensor
//  [python]
def _apply_negative_infinity_mask(tensor, mask):
  """Where mask is true, add a large negative value."""
  tensor += tf.cast(mask, tf.float32) * -INF
  tensor = tf.maximum(tensor, -INF)
  return tensor
//  [python]
def input_fn_builder(input_filepattern: str, long_seq_length: int,
                     global_seq_length: int, is_training: bool,
                     answer_encoding_method: str, drop_remainder: bool):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""
  name_to_features = get_inference_name_to_features(long_seq_length,
                                                    global_seq_length)
  if is_training:
    name_to_features.update({
        "supporting_facts": tf.FixedLenFeature([global_seq_length], tf.int64),
        "answer_types": tf.FixedLenFeature([], tf.int64)
    })
    if answer_encoding_method == "span":
      name_to_features.update({
          "answer_begins": tf.FixedLenFeature([], tf.int64),
          "answer_ends": tf.FixedLenFeature([], tf.int64),
      })
    else:
      name_to_features.update({
          "answer_bio_ids": tf.FixedLenFeature([long_seq_length], tf.int64),
      })

  def decode_record(record):
    """Decodes a record to a TensorFlow example."""
    example = tf.parse_single_example(record, name_to_features)
    for name in list(example.keys()):
      if name != "unique_ids":
        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"]
    filenames = tf.gfile.Glob(input_filepattern)
    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    if is_training:
      np.random.seed(1234)
      np.random.shuffle(filenames)
      # Assigns a subset of examples to each worker, as recommended in
      # in https://arxiv.org/pdf/1706.02677.pdf.
      tpu_context = params["context"]
      host_filenames = np.array_split(
          filenames, tpu_context.num_hosts)[tpu_context.current_host]
      host_filenames = host_filenames.tolist()
      tf.logging.info("Num input files on host %d of %d: %d",
                      tpu_context.current_host, tpu_context.num_hosts,
                      len(host_filenames))
      d = tf.data.Dataset.from_tensor_slices(host_filenames).repeat()
      # `sloppy` mode means that the interleaving is not exact. This adds
      # even more randomness to the training pipeline.
      d = d.apply(
          tf.data.experimental.parallel_interleave(
              tf.data.TFRecordDataset, sloppy=True, cycle_length=4))
      d = d.shuffle(buffer_size=batch_size * 4)
    else:
      d = tf.data.TFRecordDataset(filenames)
    d = d.map(
        map_func=decode_record,
        num_parallel_calls=tf.data.experimental.AUTOTUNE)
    d = d.batch(batch_size=batch_size, drop_remainder=drop_remainder)
    return d.prefetch(tf.data.experimental.AUTOTUNE)

  return input_fn
//  [python]
def input_fn_builder(input_filepattern: str, long_seq_length: int,
                     global_seq_length: int, is_training: bool,
                     answer_encoding_method: str, drop_remainder: bool):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""
  name_to_features = get_inference_name_to_features(long_seq_length,
                                                    global_seq_length)
  if is_training:
    name_to_features.update({
        "supporting_facts": tf.FixedLenFeature([global_seq_length], tf.int64),
        "answer_types": tf.FixedLenFeature([], tf.int64)
    })
    if answer_encoding_method == "span":
      name_to_features.update({
          "answer_begins": tf.FixedLenFeature([], tf.int64),
          "answer_ends": tf.FixedLenFeature([], tf.int64),
      })
    else:
      name_to_features.update({
          "answer_bio_ids": tf.FixedLenFeature([long_seq_length], tf.int64),
      })

  def decode_record(record):
    """Decodes a record to a TensorFlow example."""
    example = tf.parse_single_example(record, name_to_features)
    for name in list(example.keys()):
      if name != "unique_ids":
        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"]
    filenames = tf.gfile.Glob(input_filepattern)
    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    if is_training:
      np.random.seed(1234)
      np.random.shuffle(filenames)
      # Assigns a subset of examples to each worker, as recommended in
      # in https://arxiv.org/pdf/1706.02677.pdf.
      tpu_context = params["context"]
      host_filenames = np.array_split(
          filenames, tpu_context.num_hosts)[tpu_context.current_host]
      host_filenames = host_filenames.tolist()
      tf.logging.info("Num input files on host %d of %d: %d",
                      tpu_context.current_host, tpu_context.num_hosts,
                      len(host_filenames))
      d = tf.data.Dataset.from_tensor_slices(host_filenames).repeat()
      # `sloppy` mode means that the interleaving is not exact. This adds
      # even more randomness to the training pipeline.
      d = d.apply(
          tf.data.experimental.parallel_interleave(
              tf.data.TFRecordDataset, sloppy=True, cycle_length=4))
      d = d.shuffle(buffer_size=batch_size * 4)
    else:
      d = tf.data.TFRecordDataset(filenames)
    d = d.map(
        map_func=decode_record,
        num_parallel_calls=tf.data.experimental.AUTOTUNE)
    d = d.batch(batch_size=batch_size, drop_remainder=drop_remainder)
    return d.prefetch(tf.data.experimental.AUTOTUNE)

  return input_fn
//  [python]
def ssd_random_expand(image, bboxes, ratio=2., name=None):
  with tf.name_scope('ssd_random_expand'):
    image = tf.convert_to_tensor(image, name='image')
    if image.get_shape().ndims != 3:
      raise ValueError('\'image\' must have 3 dimensions.')

    height, width, depth = _ImageDimensions(image, rank=3)

    float_height, float_width = tf.to_float(height), tf.to_float(width)

    canvas_width, canvas_height = tf.to_int32(float_width * ratio), tf.to_int32(float_height * ratio)

    mean_color_of_image = [_R_MEAN/255., _G_MEAN/255., _B_MEAN/255.]#tf.reduce_mean(tf.reshape(image, [-1, 3]), 0)

    x = tf.random_uniform([], minval=0, maxval=canvas_width - width, dtype=tf.int32)
    y = tf.random_uniform([], minval=0, maxval=canvas_height - height, dtype=tf.int32)

    paddings = tf.convert_to_tensor([[y, canvas_height - height - y], [x, canvas_width - width - x]])

    big_canvas = tf.stack([tf.pad(image[:, :, 0], paddings, "CONSTANT", constant_values = mean_color_of_image[0]),
                          tf.pad(image[:, :, 1], paddings, "CONSTANT", constant_values = mean_color_of_image[1]),
                          tf.pad(image[:, :, 2], paddings, "CONSTANT", constant_values = mean_color_of_image[2])], axis=-1)

    scale = tf.cast(tf.stack([height, width, height, width]), bboxes.dtype)
    absolute_bboxes = bboxes * scale + tf.cast(tf.stack([y, x, y, x]), bboxes.dtype)

    return big_canvas, absolute_bboxes / tf.cast(tf.stack([canvas_height, canvas_width, canvas_height, canvas_width]), bboxes.dtype)
//  [python]
def ssd_random_expand(image, bboxes, ratio=2., name=None):
  with tf.name_scope('ssd_random_expand'):
    image = tf.convert_to_tensor(image, name='image')
    if image.get_shape().ndims != 3:
      raise ValueError('\'image\' must have 3 dimensions.')

    height, width, depth = _ImageDimensions(image, rank=3)

    float_height, float_width = tf.cast(height, tf.float32), tf.cast(width, tf.float32)

    canvas_width, canvas_height = tf.to_int32(float_width * ratio), tf.to_int32(float_height * ratio)

    mean_color_of_image = [_R_MEAN/255., _G_MEAN/255., _B_MEAN/255.]#tf.reduce_mean(tf.reshape(image, [-1, 3]), 0)

    x = tf.random_uniform([], minval=0, maxval=canvas_width - width, dtype=tf.int32)
    y = tf.random_uniform([], minval=0, maxval=canvas_height - height, dtype=tf.int32)

    paddings = tf.convert_to_tensor([[y, canvas_height - height - y], [x, canvas_width - width - x]])

    big_canvas = tf.stack([tf.pad(image[:, :, 0], paddings, "CONSTANT", constant_values = mean_color_of_image[0]),
                          tf.pad(image[:, :, 1], paddings, "CONSTANT", constant_values = mean_color_of_image[1]),
                          tf.pad(image[:, :, 2], paddings, "CONSTANT", constant_values = mean_color_of_image[2])], axis=-1)

    scale = tf.cast(tf.stack([height, width, height, width]), bboxes.dtype)
    absolute_bboxes = bboxes * scale + tf.cast(tf.stack([y, x, y, x]), bboxes.dtype)

    return big_canvas, absolute_bboxes / tf.cast(tf.stack([canvas_height, canvas_width, canvas_height, canvas_width]), bboxes.dtype)
//  [python]
def preprocess_for_eval(image, out_shape, data_format='channels_first', scope='ssd_preprocessing_eval', output_rgb=True):
  """Preprocesses the given image for evaluation.

  Args:
    image: A `Tensor` representing an image of arbitrary size.
    out_shape: The height and width of the image after preprocessing.
    data_format: The data_format of the desired output image.
  Returns:
    A preprocessed image.
  """
  with tf.name_scope(scope, 'ssd_preprocessing_eval', [image]):
    image = tf.to_float(image)
    image = tf.image.resize_images(image, out_shape, method=tf.image.ResizeMethod.BILINEAR, align_corners=False)
    image.set_shape(out_shape + [3])

    image = _mean_image_subtraction(image, [_R_MEAN, _G_MEAN, _B_MEAN])
    if not output_rgb:
      image_channels = tf.unstack(image, axis=-1, name='split_rgb')
      image = tf.stack([image_channels[2], image_channels[1], image_channels[0]], axis=-1, name='merge_bgr')
    # Image data format.
    if data_format == 'channels_first':
      image = tf.transpose(image, perm=(2, 0, 1))
    return image
//  [python]
def preprocess_for_eval(image, out_shape, data_format='channels_first', scope='ssd_preprocessing_eval', output_rgb=True):
  """Preprocesses the given image for evaluation.

  Args:
    image: A `Tensor` representing an image of arbitrary size.
    out_shape: The height and width of the image after preprocessing.
    data_format: The data_format of the desired output image.
  Returns:
    A preprocessed image.
  """
  with tf.name_scope(scope, 'ssd_preprocessing_eval', [image]):
    image = tf.cast(image, tf.float32)
    image = tf.image.resize_images(image, out_shape, method=tf.image.ResizeMethod.BILINEAR, align_corners=False)
    image.set_shape(out_shape + [3])

    image = _mean_image_subtraction(image, [_R_MEAN, _G_MEAN, _B_MEAN])
    if not output_rgb:
      image_channels = tf.unstack(image, axis=-1, name='split_rgb')
      image = tf.stack([image_channels[2], image_channels[1], image_channels[0]], axis=-1, name='merge_bgr')
    # Image data format.
    if data_format == 'channels_first':
      image = tf.transpose(image, perm=(2, 0, 1))
    return image
//  [python]
def _div_maybe_zero(total_loss, num_present):
  """Normalizes the total loss with the number of present pixels."""
  return tf.to_float(num_present > 0) * tf.math.divide(
      total_loss,
      tf.maximum(1e-5, num_present))
//  [python]
def _div_maybe_zero(total_loss, num_present):
  """Normalizes the total loss with the number of present pixels."""
  return tf.cast(num_present > 0, tf.float32) * tf.math.divide(
      total_loss,
      tf.maximum(1e-5, num_present))
//  [python]
def timepoint_classification_target(segments, segment_classes, num_classes):
  """Produces a classification target at each timepoint.

  If no segments are present at a time point, the first class is set to 1.0.
  This should be used as a background class unless segments are always present.

  Args:
    segments: a [time, num_segments] tensor that is 1.0 at indices within
      each segment and 0.0 elsewhere.
    segment_classes: a [num_segments] tensor with the class index of each
      segment.
    num_classes: the number of classes (must be >= max(segment_classes) + 1)
  Returns:
    a [time, num_classes] tensor. In the final output, more than one
    value in a row can be 1.0 if segments overlap.
  """
  num_segments = tf.shape(segments)[1]
  matrix_of_class_indices = tf.to_int32(
      segments * tf.to_float(tf.expand_dims(segment_classes, 0)))
  # First column will have one count per zero segment. Correct this to be 0
  # unless no segments are present.
  one_hot = tf.reduce_sum(tf.one_hot(matrix_of_class_indices, num_classes), 1)
  normalizer = tf.concat([
      tf.ones(shape=[1, 1], dtype=tf.float32) / tf.to_float(num_segments),
      tf.ones(shape=[1, num_classes - 1], dtype=tf.float32)
  ], 1)
  corrected_one_hot = tf.floor(one_hot * normalizer)
  return corrected_one_hot
//  [python]
def timepoint_classification_target(segments, segment_classes, num_classes):
  """Produces a classification target at each timepoint.

  If no segments are present at a time point, the first class is set to 1.0.
  This should be used as a background class unless segments are always present.

  Args:
    segments: a [time, num_segments] tensor that is 1.0 at indices within
      each segment and 0.0 elsewhere.
    segment_classes: a [num_segments] tensor with the class index of each
      segment.
    num_classes: the number of classes (must be >= max(segment_classes) + 1)
  Returns:
    a [time, num_classes] tensor. In the final output, more than one
    value in a row can be 1.0 if segments overlap.
  """
  num_segments = tf.shape(segments)[1]
  matrix_of_class_indices = tf.cast(
      segments * tf.cast(tf.expand_dims(segment_classes, 0), tf.float32), tf.int32)
  # First column will have one count per zero segment. Correct this to be 0
  # unless no segments are present.
  one_hot = tf.reduce_sum(tf.one_hot(matrix_of_class_indices, num_classes), 1)
  normalizer = tf.concat([
      tf.ones(shape=[1, 1], dtype=tf.float32) / tf.cast(num_segments, tf.float32),
      tf.ones(shape=[1, num_classes - 1], dtype=tf.float32)
  ], 1)
  corrected_one_hot = tf.floor(one_hot * normalizer)
  return corrected_one_hot
//  [python]
def one_hot_segments(start_indices, end_indices, num_samples):
  """Returns a one-hot, float matrix of segments at each timestep.

  All integers in the inclusive range of start_indices and end_indices are used.
  This allows start and end timestamps to be mapped to the same index and the
  segment will not be omitted.

  Args:
    start_indices: a 1d tensor of integer indices for the start of each
      segement.
    end_indices: a tensor of integer indices for the end of each segment.
      Must be the same shape as start_indices. Values should be >= start_indices
      but not strictly enforced.
    num_samples: the number of rows in the output. Indices should be <
      num_samples, but this is not strictly enforced.
  Returns:
    (segments, indicator)
    segments: A [num_samples, num_elements(start_indices)] tensor where in each
      column the rows with indices >= start_indices[column] and
      <= end_indices[column] are 1.0 and all other values are 0.0.
    indicator: a tensor of 1.0 values with shape [num_samples, 1]. If padded
      with zeros to align sizes, the indicator marks where segments is valid.
  """
  start_indices = tf.convert_to_tensor(start_indices)
  end_indices = tf.convert_to_tensor(end_indices)
  start_indices.shape.assert_is_compatible_with(end_indices.shape)
  start_indices.shape.assert_has_rank(1)
  end_indices.shape.assert_has_rank(1)
  # create a matrix of the index at each row with a column per segment.
  indices = tf.to_int64(
      tf.tile(
          tf.transpose(tf.expand_dims(tf.range(num_samples), 0)),
          [1, tf.shape(start_indices)[0]]))
  # switch to one hot encoding of segments (includes start and end indices)
  segments = tf.to_float(
      tf.logical_and(
          tf.greater_equal(indices, start_indices),
          tf.less_equal(indices, end_indices)))
  # create a tensors of ones everywhere there's an annotation. If padded with
  # zeros later, element-wise multiplication of the loss will mask out the
  # padding.
  indicator = tf.ones(shape=[num_samples, 1], dtype=tf.float32)
  return segments, indicator
//  [python]
def one_hot_segments(start_indices, end_indices, num_samples):
  """Returns a one-hot, float matrix of segments at each timestep.

  All integers in the inclusive range of start_indices and end_indices are used.
  This allows start and end timestamps to be mapped to the same index and the
  segment will not be omitted.

  Args:
    start_indices: a 1d tensor of integer indices for the start of each
      segement.
    end_indices: a tensor of integer indices for the end of each segment.
      Must be the same shape as start_indices. Values should be >= start_indices
      but not strictly enforced.
    num_samples: the number of rows in the output. Indices should be <
      num_samples, but this is not strictly enforced.
  Returns:
    (segments, indicator)
    segments: A [num_samples, num_elements(start_indices)] tensor where in each
      column the rows with indices >= start_indices[column] and
      <= end_indices[column] are 1.0 and all other values are 0.0.
    indicator: a tensor of 1.0 values with shape [num_samples, 1]. If padded
      with zeros to align sizes, the indicator marks where segments is valid.
  """
  start_indices = tf.convert_to_tensor(start_indices)
  end_indices = tf.convert_to_tensor(end_indices)
  start_indices.shape.assert_is_compatible_with(end_indices.shape)
  start_indices.shape.assert_has_rank(1)
  end_indices.shape.assert_has_rank(1)
  # create a matrix of the index at each row with a column per segment.
  indices = tf.cast(
      tf.tile(
          tf.transpose(tf.expand_dims(tf.range(num_samples), 0)),
          [1, tf.shape(start_indices)[0]]), tf.int64)
  # switch to one hot encoding of segments (includes start and end indices)
  segments = tf.to_float(
      tf.logical_and(
          tf.greater_equal(indices, start_indices),
          tf.less_equal(indices, end_indices)))
  # create a tensors of ones everywhere there's an annotation. If padded with
  # zeros later, element-wise multiplication of the loss will mask out the
  # padding.
  indicator = tf.ones(shape=[num_samples, 1], dtype=tf.float32)
  return segments, indicator
//  [python]
def non_max_suppression(scores_in,
                        boxes_in,
                        top_k_indices,
                        labels,
                        num_detections=ssd_constants.MAX_NUM_EVAL_BOXES):
  """Implement Non-maximum suppression.

  Args:
    scores_in: a Tensor with shape [batch_size,
      ssd_constants.MAX_NUM_EVAL_BOXES, num_classes]. The top
      ssd_constants.MAX_NUM_EVAL_BOXES box scores for each class.
    boxes_in: a Tensor with shape [batch_size, N, 4], which stacks box
      regression outputs on all feature levels. The N is the number of total
      anchors on all levels.
    top_k_indices: a Tensor with shape [batch_size,
      ssd_constants.MAX_NUM_EVAL_BOXES, num_classes]. The indices for these top
      boxes for each class.
    labels: labels tensor.
    num_detections: maximum output length.

  Returns:
    A tensor size of [batch_size, num_detections, 6] represents boxes, labels
    and scores after NMS.
  """

  _, _, num_classes = scores_in.get_shape().as_list()
  source_id = tf.to_float(
      tf.tile(
          tf.expand_dims(labels[ssd_constants.SOURCE_ID], 1),
          [1, num_detections]))
  raw_shape = tf.to_float(
      tf.tile(
          tf.expand_dims(labels[ssd_constants.RAW_SHAPE], 1),
          [1, num_detections, 1]))

  list_of_all_boxes = []
  list_of_all_scores = []
  list_of_all_classes = []
  # Skip background class.
  for class_i in range(1, num_classes, 1):
    boxes = tf.batch_gather(boxes_in, top_k_indices[:, :, class_i])
    class_i_scores = scores_in[:, :, class_i]
    class_i_scores, boxes = _filter_scores(class_i_scores, boxes)
    (class_i_post_scores,
     class_i_post_boxes) = ssd_architecture.non_max_suppression_padded(
         scores=tf.to_float(class_i_scores),
         boxes=tf.to_float(boxes),
         max_output_size=num_detections,
         iou_threshold=ssd_constants.OVERLAP_CRITERIA)
    class_i_classes = tf.fill(
        tf.shape(class_i_post_scores), ssd_constants.CLASS_INV_MAP[class_i])
    list_of_all_boxes.append(class_i_post_boxes)
    list_of_all_scores.append(class_i_post_scores)
    list_of_all_classes.append(class_i_classes)

  post_nms_boxes = tf.concat(list_of_all_boxes, axis=1)
  post_nms_scores = tf.concat(list_of_all_scores, axis=1)
  post_nms_classes = tf.concat(list_of_all_classes, axis=1)

  # sort all results.
  post_nms_scores, sorted_indices = tf.nn.top_k(
      tf.to_float(post_nms_scores), k=num_detections, sorted=True)

  post_nms_boxes = tf.batch_gather(post_nms_boxes, sorted_indices)
  post_nms_classes = tf.batch_gather(post_nms_classes, sorted_indices)
  detections_result = tf.stack([
      source_id,
      post_nms_boxes[:, :, 1] * raw_shape[:, :, 1],
      post_nms_boxes[:, :, 0] * raw_shape[:, :, 0],
      (post_nms_boxes[:, :, 3] - post_nms_boxes[:, :, 1]) * raw_shape[:, :, 1],
      (post_nms_boxes[:, :, 2] - post_nms_boxes[:, :, 0]) * raw_shape[:, :, 0],
      post_nms_scores,
      tf.to_float(post_nms_classes),
  ],
                               axis=2)

  return detections_result
//  [python]
def non_max_suppression(scores_in,
                        boxes_in,
                        top_k_indices,
                        labels,
                        num_detections=ssd_constants.MAX_NUM_EVAL_BOXES):
  """Implement Non-maximum suppression.

  Args:
    scores_in: a Tensor with shape [batch_size,
      ssd_constants.MAX_NUM_EVAL_BOXES, num_classes]. The top
      ssd_constants.MAX_NUM_EVAL_BOXES box scores for each class.
    boxes_in: a Tensor with shape [batch_size, N, 4], which stacks box
      regression outputs on all feature levels. The N is the number of total
      anchors on all levels.
    top_k_indices: a Tensor with shape [batch_size,
      ssd_constants.MAX_NUM_EVAL_BOXES, num_classes]. The indices for these top
      boxes for each class.
    labels: labels tensor.
    num_detections: maximum output length.

  Returns:
    A tensor size of [batch_size, num_detections, 6] represents boxes, labels
    and scores after NMS.
  """

  _, _, num_classes = scores_in.get_shape().as_list()
  source_id = tf.cast(
      tf.tile(
          tf.expand_dims(labels[ssd_constants.SOURCE_ID], 1),
          [1, num_detections]), tf.float32)
  raw_shape = tf.cast(
      tf.tile(
          tf.expand_dims(labels[ssd_constants.RAW_SHAPE], 1),
          [1, num_detections, 1]), tf.float32)

  list_of_all_boxes = []
  list_of_all_scores = []
  list_of_all_classes = []
  # Skip background class.
  for class_i in range(1, num_classes, 1):
    boxes = tf.batch_gather(boxes_in, top_k_indices[:, :, class_i])
    class_i_scores = scores_in[:, :, class_i]
    class_i_scores, boxes = _filter_scores(class_i_scores, boxes)
    (class_i_post_scores,
     class_i_post_boxes) = ssd_architecture.non_max_suppression_padded(
         scores=tf.cast(class_i_scores, tf.float32),
         boxes=tf.cast(boxes, tf.float32),
         max_output_size=num_detections,
         iou_threshold=ssd_constants.OVERLAP_CRITERIA)
    class_i_classes = tf.fill(
        tf.shape(class_i_post_scores), ssd_constants.CLASS_INV_MAP[class_i])
    list_of_all_boxes.append(class_i_post_boxes)
    list_of_all_scores.append(class_i_post_scores)
    list_of_all_classes.append(class_i_classes)

  post_nms_boxes = tf.concat(list_of_all_boxes, axis=1)
  post_nms_scores = tf.concat(list_of_all_scores, axis=1)
  post_nms_classes = tf.concat(list_of_all_classes, axis=1)

  # sort all results.
  post_nms_scores, sorted_indices = tf.nn.top_k(
      tf.cast(post_nms_scores, tf.float32), k=num_detections, sorted=True)

  post_nms_boxes = tf.batch_gather(post_nms_boxes, sorted_indices)
  post_nms_classes = tf.batch_gather(post_nms_classes, sorted_indices)
  detections_result = tf.stack([
      source_id,
      post_nms_boxes[:, :, 1] * raw_shape[:, :, 1],
      post_nms_boxes[:, :, 0] * raw_shape[:, :, 0],
      (post_nms_boxes[:, :, 3] - post_nms_boxes[:, :, 1]) * raw_shape[:, :, 1],
      (post_nms_boxes[:, :, 2] - post_nms_boxes[:, :, 0]) * raw_shape[:, :, 0],
      post_nms_scores,
      tf.cast(post_nms_classes, tf.float32),
  ],
                               axis=2)

  return detections_result
//  [python]
def _apply_drop_path(self, x, layer_id):
    drop_path_keep_prob = self.drop_path_keep_prob

    layer_ratio = float(layer_id + 1) / (self.num_layers + 2)
    drop_path_keep_prob = 1.0 - layer_ratio * (1.0 - drop_path_keep_prob)

    step_ratio = tf.to_float(self.global_step + 1) / tf.to_float(self.num_train_steps)
    step_ratio = tf.minimum(1.0, step_ratio)
    drop_path_keep_prob = 1.0 - step_ratio * (1.0 - drop_path_keep_prob)

    x = drop_path(x, drop_path_keep_prob)
    return x
//  [python]
def _apply_drop_path(self, x, layer_id):
    drop_path_keep_prob = self.drop_path_keep_prob

    layer_ratio = float(layer_id + 1) / (self.num_layers + 2)
    drop_path_keep_prob = 1.0 - layer_ratio * (1.0 - drop_path_keep_prob)

    step_ratio = tf.cast(self.global_step + 1, tf.float32) / tf.cast(self.num_train_steps, tf.float32)
    step_ratio = tf.minimum(1.0, step_ratio)
    drop_path_keep_prob = 1.0 - step_ratio * (1.0 - drop_path_keep_prob)

    x = drop_path(x, drop_path_keep_prob)
    return x
//  [python]
def _build_valid(self):
    if self.x_valid is not None:
      print("-" * 80)
      print("Build valid graph")
      logits = self._model(self.x_valid, False, reuse=True)
      self.valid_preds = tf.argmax(logits, axis=1)
      self.valid_preds = tf.to_int32(self.valid_preds)
      self.valid_acc = tf.equal(self.valid_preds, self.y_valid)
      self.valid_acc = tf.to_int32(self.valid_acc)
      self.valid_acc = tf.reduce_sum(self.valid_acc)
//  [python]
def _build_valid(self):
    if self.x_valid is not None:
      print("-" * 80)
      print("Build valid graph")
      logits = self._model(self.x_valid, False, reuse=True)
      self.valid_preds = tf.argmax(logits, axis=1)
      self.valid_preds = tf.cast(self.valid_preds, tf.int32)
      self.valid_acc = tf.equal(self.valid_preds, self.y_valid)
      self.valid_acc = tf.cast(self.valid_acc, tf.int32)
      self.valid_acc = tf.reduce_sum(self.valid_acc)
//  [python]
def _build_test(self):
    print("-" * 80)
    print("Build test graph")
    logits = self._model(self.x_test, False, reuse=True)
    self.test_preds = tf.argmax(logits, axis=1)
    self.test_preds = tf.to_int32(self.test_preds)
    self.test_acc = tf.equal(self.test_preds, self.y_test)
    self.test_acc = tf.to_int32(self.test_acc)
    self.test_acc = tf.reduce_sum(self.test_acc)
//  [python]
def _build_test(self):
    print("-" * 80)
    print("Build test graph")
    logits = self._model(self.x_test, False, reuse=True)
    self.test_preds = tf.argmax(logits, axis=1)
    self.test_preds = tf.cast(self.test_preds, tf.int32)
    self.test_acc = tf.equal(self.test_preds, self.y_test)
    self.test_acc = tf.cast(self.test_acc, tf.int32)
    self.test_acc = tf.reduce_sum(self.test_acc)
//  [python]
def reduce_logavgexp(input_tensor, axis=None, keepdims=None, name=None):
  dims = tf.shape(input_tensor)
  if axis is not None:
    dims = tf.gather(dims, axis)
  denominator = tf.reduce_prod(dims)
  return (tf.reduce_logsumexp(
      input_tensor, axis=axis, keepdims=keepdims, name=name) -
          tf.log(tf.to_float(denominator)))
//  [python]
def reduce_logavgexp(input_tensor, axis=None, keepdims=None, name=None):
  dims = tf.shape(input_tensor)
  if axis is not None:
    dims = tf.gather(dims, axis)
  denominator = tf.reduce_prod(dims)
  return (tf.reduce_logsumexp(
      input_tensor, axis=axis, keepdims=keepdims, name=name) -
          tf.log(tf.cast(denominator, tf.float32)))
//  [python]
def add_intrinsics_head(bottleneck, image_height, image_width):
  """Adds a head the preficts camera intrinsics.

  Args:
    bottleneck: A tf.Tensor of shape [B, 1, 1, C], typically the bottlenech
      features of a netrowk.
    image_height: A scalar tf.Tensor or an python scalar, the image height in
      pixels.
    image_width: A scalar tf.Tensor or an python scalar, the image width in
      pixels.

  image_height and image_width are used to provide the right scale for the focal
  length and the offest parameters.

  Returns:
    a tf.Tensor of shape [B, 3, 3], and type float32, where the 3x3 part is the
    intrinsic matrix: (fx, 0, x0), (0, fy, y0), (0, 0, 1).
  """
  with tf.variable_scope('CameraIntrinsics'):
    # Since the focal lengths in pixels tend to be in the order of magnitude of
    # the image width and height, we multiply the network prediction by them.
    focal_lengths = tf.squeeze(
        layers.conv2d(
            bottleneck,
            2, [1, 1],
            stride=1,
            activation_fn=tf.nn.softplus,
            weights_regularizer=None,
            scope='foci'),
        axis=(1, 2)) * tf.to_float(
            tf.convert_to_tensor([[image_width, image_height]]))

    # The pixel offsets tend to be around the center of the image, and they
    # are typically a fraction the image width and height in pixels. We thus
    # multiply the network prediction by the width and height, and the
    # additional 0.5 them by default at the center of the image.
    offsets = (tf.squeeze(
        layers.conv2d(
            bottleneck,
            2, [1, 1],
            stride=1,
            activation_fn=None,
            weights_regularizer=None,
            biases_initializer=None,
            scope='offsets'),
        axis=(1, 2)) + 0.5) * tf.to_float(
            tf.convert_to_tensor([[image_width, image_height]]))

    foci = tf.linalg.diag(focal_lengths)

    maybe_summary.scalar('foci', tf.reduce_mean(foci))
    maybe_summary.scalar('offsets', tf.reduce_mean(offsets))

    intrinsic_mat = tf.concat([foci, tf.expand_dims(offsets, -1)], axis=2)
    batch_size = tf.shape(bottleneck)[0]
    last_row = tf.tile([[[0.0, 0.0, 1.0]]], [batch_size, 1, 1])
    intrinsic_mat = tf.concat([intrinsic_mat, last_row], axis=1)
    return intrinsic_mat
//  [python]
def add_intrinsics_head(bottleneck, image_height, image_width):
  """Adds a head the preficts camera intrinsics.

  Args:
    bottleneck: A tf.Tensor of shape [B, 1, 1, C], typically the bottlenech
      features of a netrowk.
    image_height: A scalar tf.Tensor or an python scalar, the image height in
      pixels.
    image_width: A scalar tf.Tensor or an python scalar, the image width in
      pixels.

  image_height and image_width are used to provide the right scale for the focal
  length and the offest parameters.

  Returns:
    a tf.Tensor of shape [B, 3, 3], and type float32, where the 3x3 part is the
    intrinsic matrix: (fx, 0, x0), (0, fy, y0), (0, 0, 1).
  """
  with tf.variable_scope('CameraIntrinsics'):
    # Since the focal lengths in pixels tend to be in the order of magnitude of
    # the image width and height, we multiply the network prediction by them.
    focal_lengths = tf.squeeze(
        layers.conv2d(
            bottleneck,
            2, [1, 1],
            stride=1,
            activation_fn=tf.nn.softplus,
            weights_regularizer=None,
            scope='foci'),
        axis=(1, 2)) * tf.cast(
            tf.convert_to_tensor([[image_width, image_height]]), tf.float32)

    # The pixel offsets tend to be around the center of the image, and they
    # are typically a fraction the image width and height in pixels. We thus
    # multiply the network prediction by the width and height, and the
    # additional 0.5 them by default at the center of the image.
    offsets = (tf.squeeze(
        layers.conv2d(
            bottleneck,
            2, [1, 1],
            stride=1,
            activation_fn=None,
            weights_regularizer=None,
            biases_initializer=None,
            scope='offsets'),
        axis=(1, 2)) + 0.5) * tf.cast(
            tf.convert_to_tensor([[image_width, image_height]]), tf.float32)

    foci = tf.linalg.diag(focal_lengths)

    maybe_summary.scalar('foci', tf.reduce_mean(foci))
    maybe_summary.scalar('offsets', tf.reduce_mean(offsets))

    intrinsic_mat = tf.concat([foci, tf.expand_dims(offsets, -1)], axis=2)
    batch_size = tf.shape(bottleneck)[0]
    last_row = tf.tile([[[0.0, 0.0, 1.0]]], [batch_size, 1, 1])
    intrinsic_mat = tf.concat([intrinsic_mat, last_row], axis=1)
    return intrinsic_mat
//  [python]
def _build_train(self):
    print("-" * 80)
    print("Build train graph")
    logits = self._model(self.x_train, is_training=True)
    log_probs = tf.nn.sparse_softmax_cross_entropy_with_logits(
      logits=logits, labels=self.y_train)
    self.loss = tf.reduce_mean(log_probs)

    if self.use_aux_heads:
      log_probs = tf.nn.sparse_softmax_cross_entropy_with_logits(
        logits=self.aux_logits, labels=self.y_train)
      self.aux_loss = tf.reduce_mean(log_probs)
      train_loss = self.loss + 0.4 * self.aux_loss
    else:
      train_loss = self.loss

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

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

    self.train_op, self.lr, self.grad_norm, self.optimizer = get_train_ops(
      train_loss,
      tf_variables,
      self.global_step,
      clip_mode=self.clip_mode,
      grad_bound=self.grad_bound,
      l2_reg=self.l2_reg,
      lr_init=self.lr_init,
      lr_dec_start=self.lr_dec_start,
      lr_dec_every=self.lr_dec_every,
      lr_dec_rate=self.lr_dec_rate,
      lr_cosine=self.lr_cosine,
      lr_max=self.lr_max,
      lr_min=self.lr_min,
      lr_T_0=self.lr_T_0,
      lr_T_mul=self.lr_T_mul,
      num_train_batches=self.num_train_batches,
      optim_algo=self.optim_algo,
      sync_replicas=self.sync_replicas,
      num_aggregate=self.num_aggregate,
      num_replicas=self.num_replicas)