jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def add_embeddings(self):
        with tf.variable_scope('embedding'):
            if self.config.embeddings is not None:
                embeddings = tf.Variable(self.config.embeddings, 
                    name="embeddings", trainable=False)
            else:
                embeddings = tf.get_variable('embeddings', 
                    shape=[self.config.vocab_size, self.config.embedding_size], 
                    initializer=tf.uniform_unit_scaling_initializer())
            q_embed = tf.nn.embedding_lookup(embeddings, self.q)
            a_embed = tf.nn.embedding_lookup(embeddings, self.a)
            return q_embed, a_embed
//  [python]
def add_embeddings(self):
        with tf.variable_scope('embedding'):
            if self.config.embeddings is not None:
                embeddings = tf.Variable(self.config.embeddings, 
                    name="embeddings", trainable=False)
            else:
                embeddings = tf.get_variable('embeddings', 
                    shape=[self.config.vocab_size, self.config.embedding_size], 
                    initializer=tf.keras.initializers.VarianceScaling())
            q_embed = tf.nn.embedding_lookup(embeddings, self.q)
            a_embed = tf.nn.embedding_lookup(embeddings, self.a)
            return q_embed, a_embed
//  [python]
def _fully_connected(self, x, out_dim):
    """FullyConnected layer for final output."""
    num_non_batch_dimensions = len(x.shape)
    prod_non_batch_dimensions = 1
    for ii in range(num_non_batch_dimensions - 1):
      prod_non_batch_dimensions *= int(x.shape[ii + 1])
    x = tf.reshape(x, [tf.shape(x)[0], -1])
    w = tf.get_variable(
        'DW', [prod_non_batch_dimensions, out_dim],
        initializer=tf.uniform_unit_scaling_initializer(factor=1.0))
    b = tf.get_variable('biases', [out_dim],
                        initializer=tf.constant_initializer())
    return tf.nn.xw_plus_b(x, w, b)
//  [python]
def _fully_connected(self, x, out_dim):
    """FullyConnected layer for final output."""
    num_non_batch_dimensions = len(x.shape)
    prod_non_batch_dimensions = 1
    for ii in range(num_non_batch_dimensions - 1):
      prod_non_batch_dimensions *= int(x.shape[ii + 1])
    x = tf.reshape(x, [tf.shape(x)[0], -1])
    w = tf.get_variable(
        'DW', [prod_non_batch_dimensions, out_dim],
        initializer=tf.keras.initializers.VarianceScaling(scale=1.0))
    b = tf.get_variable('biases', [out_dim],
                        initializer=tf.constant_initializer())
    return tf.nn.xw_plus_b(x, w, b)
//  [python]
def _fully_connected(self, x, out_dim):
    x = tf.reshape(x, [self._params.batch_size, -1])
    w = tf.get_variable(
        'DW', [x.get_shape()[1], out_dim],
        initializer=tf.uniform_unit_scaling_initializer(factor=1.0))
    b = tf.get_variable(
        'biases', [out_dim], initializer=tf.constant_initializer())
    return tf.nn.xw_plus_b(x, w, b)
//  [python]
def _fully_connected(self, x, out_dim):
    x = tf.reshape(x, [self._params.batch_size, -1])
    w = tf.get_variable(
        'DW', [x.get_shape()[1], out_dim],
        initializer=tf.initializers.variance_scaling(factor=1.0))
    b = tf.get_variable(
        'biases', [out_dim], initializer=tf.constant_initializer())
    return tf.nn.xw_plus_b(x, w, b)
//  [python]
def conv1d(
    x,
    num_filters,
    filter_length,
    name,
    dilation=1,
    causal=True,
    kernel_initializer=tf.uniform_unit_scaling_initializer(1.0),
    biases_initializer=tf.constant_initializer(0.0),
    is_training=True,
):
    """Fast 1D convolution that supports causal padding and dilation.

  Args:
    x: The [mb, time, channels] float tensor that we convolve.
    num_filters: The number of filter maps in the convolution.
    filter_length: The integer length of the filter.
    name: The name of the scope for the variables.
    dilation: The amount of dilation.
    causal: Whether or not this is a causal convolution.
    kernel_initializer: The kernel initialization function.
    biases_initializer: The biases initialization function.
    is_training: Whether or not ot use traininable variables.

  Returns:
    y: The output of the 1D convolution.
  """
    batch_size, length, num_input_channels = x.get_shape().as_list()
    assert length % dilation == 0

    kernel_shape = [1, filter_length, num_input_channels, num_filters]
    strides = [1, 1, 1, 1]
    biases_shape = [num_filters]
    padding = 'VALID' if causal else 'SAME'

    with tf.variable_scope(name):
        weights = tf.get_variable(
            'W',
            shape=kernel_shape,
            initializer=kernel_initializer,
            trainable=is_training,
        )
        biases = tf.get_variable(
            'biases',
            shape=biases_shape,
            initializer=biases_initializer,
            trainable=is_training,
        )

    x_ttb = time_to_batch(x, dilation)
    if filter_length > 1 and causal:
        x_ttb = tf.pad(x_ttb, [[0, 0], [filter_length - 1, 0], [0, 0]])

    x_ttb_shape = x_ttb.get_shape().as_list()
    x_4d = tf.reshape(
        x_ttb, [x_ttb_shape[0], 1, x_ttb_shape[1], num_input_channels]
    )
    y = tf.nn.conv2d(x_4d, weights, strides, padding=padding)
    y = tf.nn.bias_add(y, biases)
    y_shape = y.get_shape().as_list()
    y = tf.reshape(y, [y_shape[0], y_shape[2], num_filters])
    y = batch_to_time(y, dilation)
    y.set_shape([batch_size, length, num_filters])
    return y
//  [python]
def conv1d(
    x,
    num_filters,
    filter_length,
    name,
    dilation=1,
    causal=True,
    kernel_initializer=tf.initializers.variance_scaling(1.0),
    biases_initializer=tf.constant_initializer(0.0),
    is_training=True,
):
    """Fast 1D convolution that supports causal padding and dilation.

  Args:
    x: The [mb, time, channels] float tensor that we convolve.
    num_filters: The number of filter maps in the convolution.
    filter_length: The integer length of the filter.
    name: The name of the scope for the variables.
    dilation: The amount of dilation.
    causal: Whether or not this is a causal convolution.
    kernel_initializer: The kernel initialization function.
    biases_initializer: The biases initialization function.
    is_training: Whether or not ot use traininable variables.

  Returns:
    y: The output of the 1D convolution.
  """
    batch_size, length, num_input_channels = x.get_shape().as_list()
    assert length % dilation == 0

    kernel_shape = [1, filter_length, num_input_channels, num_filters]
    strides = [1, 1, 1, 1]
    biases_shape = [num_filters]
    padding = 'VALID' if causal else 'SAME'

    with tf.variable_scope(name):
        weights = tf.get_variable(
            'W',
            shape=kernel_shape,
            initializer=kernel_initializer,
            trainable=is_training,
        )
        biases = tf.get_variable(
            'biases',
            shape=biases_shape,
            initializer=biases_initializer,
            trainable=is_training,
        )

    x_ttb = time_to_batch(x, dilation)
    if filter_length > 1 and causal:
        x_ttb = tf.pad(x_ttb, [[0, 0], [filter_length - 1, 0], [0, 0]])

    x_ttb_shape = x_ttb.get_shape().as_list()
    x_4d = tf.reshape(
        x_ttb, [x_ttb_shape[0], 1, x_ttb_shape[1], num_input_channels]
    )
    y = tf.nn.conv2d(x_4d, weights, strides, padding=padding)
    y = tf.nn.bias_add(y, biases)
    y_shape = y.get_shape().as_list()
    y = tf.reshape(y, [y_shape[0], y_shape[2], num_filters])
    y = batch_to_time(y, dilation)
    y.set_shape([batch_size, length, num_filters])
    return y
//  [python]
def get_vq_codebook(codebook_size, hidden_size):
  """Get lookup table for VQ bottleneck."""
  with tf.variable_scope("vq", reuse=tf.AUTO_REUSE):
    means = tf.get_variable(
        name="means",
        shape=[codebook_size, hidden_size],
        initializer=tf.uniform_unit_scaling_initializer())

    ema_count = tf.get_variable(
        name="ema_count",
        shape=[codebook_size],
        initializer=tf.constant_initializer(0),
        trainable=False)

    with tf.colocate_with(means):
      ema_means = tf.get_variable(
          name="ema_means",
          initializer=means.initialized_value(),
          trainable=False)

  return means, ema_means, ema_count
//  [python]
def get_vq_codebook(codebook_size, hidden_size):
  """Get lookup table for VQ bottleneck."""
  with tf.variable_scope("vq", reuse=tf.AUTO_REUSE):
    means = tf.get_variable(
        name="means",
        shape=[codebook_size, hidden_size],
        initializer=tf.initializers.variance_scaling())

    ema_count = tf.get_variable(
        name="ema_count",
        shape=[codebook_size],
        initializer=tf.constant_initializer(0),
        trainable=False)

    with tf.colocate_with(means):
      ema_means = tf.get_variable(
          name="ema_means",
          initializer=means.initialized_value(),
          trainable=False)

  return means, ema_means, ema_count
//  [python]
def add_embeddings(self):
        """
        embedding层
        """
        with tf.variable_scope('embedding'):
            if self.config.embeddings is not None:
                embeddings = tf.Variable(self.config.embeddings, name="embeddings", trainable=False)
            else:
                embeddings = tf.get_variable('embeddings', shape=[self.config.vocab_size, self.config.embedding_size], initializer=tf.uniform_unit_scaling_initializer())
            q_embed = tf.nn.embedding_lookup(embeddings, self.q)
            a_embed = tf.nn.embedding_lookup(embeddings, self.a)
            return q_embed, a_embed
//  [python]
def add_embeddings(self):
        """
        embedding层
        """
        with tf.variable_scope('embedding'):
            if self.config.embeddings is not None:
                embeddings = tf.Variable(self.config.embeddings, name="embeddings", trainable=False)
            else:
                embeddings = tf.get_variable('embeddings', shape=[self.config.vocab_size, self.config.embedding_size], initializer=tf.keras.initializers.VarianceScaling())
            q_embed = tf.nn.embedding_lookup(embeddings, self.q)
            a_embed = tf.nn.embedding_lookup(embeddings, self.a)
            return q_embed, a_embed
//  [python]
def make_vars(self):
        with tf.variable_scope("logit", reuse=tf.AUTO_REUSE):
            w = tf.get_variable(
                "DW",
                [self.dim, self.num_hid],
                initializer=tf.contrib.layers.variance_scaling_initializer(distribution="uniform"),
            )
            b = tf.get_variable(
                "biases", [self.num_hid], initializer=tf.zeros_initializer()
            )
        return w, b
//  [python]
def make_vars(self):
        with tf.variable_scope("logit", reuse=tf.AUTO_REUSE):
            w = tf.get_variable(
                "DW",
                [self.dim, self.num_hid],
                initializer=tf.initializers.variance_scaling(distribution="uniform"),
            )
            b = tf.get_variable(
                "biases", [self.num_hid], initializer=tf.initializers.constant()
            )
        return w, b
//  [python]
def tearDown(self):
    tf.config.experimental_run_functions_eagerly(False)
    super(PpoLearnerTest, self).tearDown()
//  [python]
def tearDown(self):
    tf.config.run_functions_eagerly(False)
    super(PpoLearnerTest, self).tearDown()
//  [python]
def setUp(self):
    super(PpoLearnerTest, self).setUp()
    # Run in full eager mode in order to inspect the content of the trajectory
    # fed into the fake agent.
    tf.config.experimental_run_functions_eagerly(True)
//  [python]
def setUp(self):
    super(PpoLearnerTest, self).setUp()
    # Run in full eager mode in order to inspect the content of the trajectory
    # fed into the fake agent.
    tf.config.run_functions_eagerly(True)
//  [python]
def testTransformerAutoencoder(self):
    hparams = imagetransformer_latent_tiny()
    hparams.mode = tf.estimator.ModeKeys.TRAIN
    block_dim = int(hparams.hidden_size // hparams.num_blocks)
    block_v_size = 2**(hparams.bottleneck_bits /
                       (hparams.num_residuals * hparams.num_blocks))
    block_v_size = int(block_v_size)
    means = tf.get_variable(
        name="means",
        shape=[hparams.num_residuals,
               hparams.num_blocks,
               block_v_size,
               block_dim],
        initializer=tf.uniform_unit_scaling_initializer())
    hparams.bottleneck = functools.partial(
        discretization.discrete_bottleneck,
        hidden_size=hparams.hidden_size,
        z_size=hparams.bottleneck_bits,
        filter_size=hparams.filter_size,
        startup_steps=hparams.startup_steps,
        bottleneck_kind=hparams.bottleneck_kind,
        num_blocks=hparams.num_blocks,
        num_residuals=hparams.num_residuals,
        reshape_method=hparams.reshape_method,
        beta=hparams.vq_beta,
        decay=hparams.vq_decay,
        soft_em=hparams.soft_em,
        num_samples=hparams.num_samples,
        epsilon=hparams.vq_epsilon,
        ema=hparams.ema,
        means=means)

    inputs = None
    batch_size = hparams.batch_size
    targets = tf.random_uniform([batch_size,
                                 hparams.img_len,
                                 hparams.img_len,
                                 hparams.hidden_size],
                                minval=-1., maxval=1.)
    target_space_id = None

    tf.train.create_global_step()
    decoder_output, losses, cache = latent_layers.transformer_autoencoder(
        inputs, targets, target_space_id, hparams)

    self.assertEqual(set(losses), {"extra", "extra_loss", "latent_pred"})

    self.evaluate(tf.global_variables_initializer())
    decoder_output_, extra_loss_, latent_pred_ = self.evaluate(
        [decoder_output, losses["extra_loss"], losses["latent_pred"]])
    self.assertEqual(decoder_output_.shape, (batch_size,
                                             hparams.img_len,
                                             hparams.img_len,
                                             hparams.hidden_size))
    self.assertEqual(extra_loss_.shape, (batch_size,))
    self.assertEqual(latent_pred_.shape, (batch_size,))
    self.assertAllGreaterEqual(extra_loss_, 0.)
    self.assertAllGreaterEqual(latent_pred_, 0.)
    self.assertEqual(cache, None)
//  [python]
def testTransformerAutoencoder(self):
    hparams = imagetransformer_latent_tiny()
    hparams.mode = tf.estimator.ModeKeys.TRAIN
    block_dim = int(hparams.hidden_size // hparams.num_blocks)
    block_v_size = 2**(hparams.bottleneck_bits /
                       (hparams.num_residuals * hparams.num_blocks))
    block_v_size = int(block_v_size)
    means = tf.get_variable(
        name="means",
        shape=[hparams.num_residuals,
               hparams.num_blocks,
               block_v_size,
               block_dim],
        initializer=tf.initializers.variance_scaling())
    hparams.bottleneck = functools.partial(
        discretization.discrete_bottleneck,
        hidden_size=hparams.hidden_size,
        z_size=hparams.bottleneck_bits,
        filter_size=hparams.filter_size,
        startup_steps=hparams.startup_steps,
        bottleneck_kind=hparams.bottleneck_kind,
        num_blocks=hparams.num_blocks,
        num_residuals=hparams.num_residuals,
        reshape_method=hparams.reshape_method,
        beta=hparams.vq_beta,
        decay=hparams.vq_decay,
        soft_em=hparams.soft_em,
        num_samples=hparams.num_samples,
        epsilon=hparams.vq_epsilon,
        ema=hparams.ema,
        means=means)

    inputs = None
    batch_size = hparams.batch_size
    targets = tf.random_uniform([batch_size,
                                 hparams.img_len,
                                 hparams.img_len,
                                 hparams.hidden_size],
                                minval=-1., maxval=1.)
    target_space_id = None

    tf.train.create_global_step()
    decoder_output, losses, cache = latent_layers.transformer_autoencoder(
        inputs, targets, target_space_id, hparams)

    self.assertEqual(set(losses), {"extra", "extra_loss", "latent_pred"})

    self.evaluate(tf.global_variables_initializer())
    decoder_output_, extra_loss_, latent_pred_ = self.evaluate(
        [decoder_output, losses["extra_loss"], losses["latent_pred"]])
    self.assertEqual(decoder_output_.shape, (batch_size,
                                             hparams.img_len,
                                             hparams.img_len,
                                             hparams.hidden_size))
    self.assertEqual(extra_loss_.shape, (batch_size,))
    self.assertEqual(latent_pred_.shape, (batch_size,))
    self.assertAllGreaterEqual(extra_loss_, 0.)
    self.assertAllGreaterEqual(latent_pred_, 0.)
    self.assertEqual(cache, None)
//  [python]
def setUp(self):
    super(PpoActorNetworkTest, self).setUp()
    # Run in full eager mode in order to inspect the content of tensors.
    tf.config.experimental_run_functions_eagerly(True)
//  [python]
def setUp(self):
    super(PpoActorNetworkTest, self).setUp()
    tf.config.run_functions_eagerly(True)
//  [python]
def _encoder(self, x):
    '''
    Note that we need a pair of reversal to ensure causality.
    (i.e. no trailing pads)
    `x`: [b, T, c]
    '''
    k_init = self.arch['initial_filter_width']

    b = tf.shape(x)[0]
    o = tf.zeros([b, k_init - 1, self.arch['dim_symbol_emb']])
    x = tf.concat([o, x], 1)

    k_init = self.arch['initial_filter_width']
    x = tf.layers.conv1d(
      inputs=x,
      filters=self.arch['residual_channels'],
      kernel_size=k_init,
      kernel_regularizer=tf.keras.regularizers.l2(WEIGHT_DECAY),
      name='initial_filtering',
      kernel_initializer=tf.glorot_uniform,
    )
    x = tf.nn.leaky_relu(x, 2e-2)

    x = tf.reverse(x, [1])  # paired op to enforce causality
    for i in range(self.arch['n_downsample_stack']):
      conv = tf.layers.conv1d(
        inputs=x,
        filters=(i + 1) * self.arch['encoder']['filters'],
        kernel_size=self.arch['encoder']['kernels'],
        strides=2,
        padding='same',
        # activation=tf.nn.tanh,
        kernel_initializer=tf.glorot_uniform,
        kernel_regularizer=tf.keras.regularizers.l2(WEIGHT_DECAY),
      )
      gate = tf.layers.conv1d(
        inputs=x,
        filters=(i + 1) * self.arch['encoder']['filters'],
        kernel_size=self.arch['encoder']['kernels'],
        strides=2,
        padding='same',
        # activation=tf.nn.sigmoid,
        kernel_initializer=tf.glorot_uniform,
        kernel_regularizer=tf.keras.regularizers.l2(WEIGHT_DECAY),
        bias_initializer=tf.initializers.ones,
      )
      x = tf.nn.tanh(conv) * tf.nn.sigmoid(gate)
    x = tf.reverse(x, [1])  # paired op to enforce causality

    x = tf.layers.conv1d(
      inputs=x,
      filters=self.arch['dim_exemplar'],
      kernel_size=1,
      kernel_initializer=tf.glorot_uniform,
      kernel_regularizer=tf.keras.regularizers.l2(WEIGHT_DECAY),
    )
    return x
//  [python]
def _encoder(self, x):
    '''
    Note that we need a pair of reversal to ensure causality.
    (i.e. no trailing pads)
    `x`: [b, T, c]
    '''
    k_init = self.arch['initial_filter_width']

    b = tf.shape(x)[0]
    o = tf.zeros([b, k_init - 1, self.arch['dim_symbol_emb']])
    x = tf.concat([o, x], 1)

    k_init = self.arch['initial_filter_width']
    x = tf.layers.conv1d(
      inputs=x,
      filters=self.arch['residual_channels'],
      kernel_size=k_init,
      kernel_regularizer=tf.keras.regularizers.l2(WEIGHT_DECAY),
      name='initial_filtering',
      kernel_initializer=tf.initializers.variance_scaling(
        scale=1.43,
        distribution='uniform'),
    )
    x = tf.nn.leaky_relu(x, 2e-2)

    x = tf.reverse(x, [1])  # paired op to enforce causality
    for i in range(self.arch['n_downsample_stack']):
      conv = tf.layers.conv1d(
        inputs=x,
        filters=(i + 1) * self.arch['encoder']['filters'],
        kernel_size=self.arch['encoder']['kernels'],
        strides=2,
        padding='same',
        # activation=tf.nn.tanh,
        kernel_initializer=tf.initializers.variance_scaling(
          scale=1.15,
          distribution='uniform'),
        kernel_regularizer=tf.keras.regularizers.l2(WEIGHT_DECAY),
      )
      gate = tf.layers.conv1d(
        inputs=x,
        filters=(i + 1) * self.arch['encoder']['filters'],
        kernel_size=self.arch['encoder']['kernels'],
        strides=2,
        padding='same',
        # activation=tf.nn.sigmoid,
        kernel_initializer=tf.initializers.variance_scaling(distribution='uniform'),
        kernel_regularizer=tf.keras.regularizers.l2(WEIGHT_DECAY),
        bias_initializer=tf.initializers.ones,
      )
      x = tf.nn.tanh(conv) * tf.nn.sigmoid(gate)
    x = tf.reverse(x, [1])  # paired op to enforce causality

    x = tf.layers.conv1d(
      inputs=x,
      filters=self.arch['dim_exemplar'],
      kernel_size=1,
      kernel_initializer=tf.initializers.variance_scaling(
          distribution='uniform'),
      kernel_regularizer=tf.keras.regularizers.l2(WEIGHT_DECAY),
    )
    return x
//  [python]
def main(_):

  imgs = [np.array(Image.open(FLAGS.image_path))] * 2
  # Create model config.
  config = hparams_config.get_efficientdet_config('efficientdet-d0')
  config.is_training_bn = False
  config.image_size = '1920x1280'
  config.nms_configs.score_thresh = 0.4
  config.nms_configs.max_output_size = 100
  config.override(FLAGS.hparams)

  # Use 'mixed_float16' if running on GPUs.
  policy = tf.keras.mixed_precision.experimental.Policy('float32')
  tf.keras.mixed_precision.experimental.set_policy(policy)
  tf.config.experimental_run_functions_eagerly(FLAGS.debug)

  # Create and run the model.
  model = efficientdet_keras.EfficientDetModel(config=config)
  model.build((None, None, None, 3))
  model.load_weights(tf.train.latest_checkpoint(FLAGS.model_dir))
  model.summary()

  class ExportModel(tf.Module):

    def __init__(self, model):
      super().__init__()
      self.model = model

    @tf.function
    def f(self, imgs):
      return self.model(imgs, training=False, post_mode='global')

  imgs = tf.convert_to_tensor(imgs, dtype=tf.uint8)
  export_model = ExportModel(model)
  if FLAGS.saved_model_dir:
    tf.saved_model.save(
        export_model,
        FLAGS.saved_model_dir,
        signatures=export_model.f.get_concrete_function(
            tf.TensorSpec(shape=(None, None, None, 3), dtype=tf.uint8)))
    export_model = tf.saved_model.load(FLAGS.saved_model_dir)

  boxes, scores, classes, valid_len = export_model.f(imgs)

  # Visualize results.
  for i, img in enumerate(imgs):
    length = valid_len[i]
    img = inference.visualize_image(
        img,
        boxes[i].numpy()[:length],
        classes[i].numpy().astype(np.int)[:length],
        scores[i].numpy()[:length],
        label_map=config.label_map,
        min_score_thresh=config.nms_configs.score_thresh,
        max_boxes_to_draw=config.nms_configs.max_output_size)
    output_image_path = os.path.join(FLAGS.output_dir, str(i) + '.jpg')
    Image.fromarray(img).save(output_image_path)
    print('writing annotated image to %s' % output_image_path)
//  [python]
def main(_):

  imgs = [np.array(Image.open(FLAGS.image_path))] * 2
  # Create model config.
  config = hparams_config.get_efficientdet_config('efficientdet-d0')
  config.is_training_bn = False
  config.image_size = '1920x1280'
  config.nms_configs.score_thresh = 0.4
  config.nms_configs.max_output_size = 100
  config.override(FLAGS.hparams)

  # Use 'mixed_float16' if running on GPUs.
  policy = tf.keras.mixed_precision.experimental.Policy('float32')
  tf.keras.mixed_precision.experimental.set_policy(policy)
  tf.config.run_functions_eagerly(FLAGS.debug)

  # Create and run the model.
  model = efficientdet_keras.EfficientDetModel(config=config)
  model.build((None, None, None, 3))
  model.load_weights(tf.train.latest_checkpoint(FLAGS.model_dir))
  model.summary()

  class ExportModel(tf.Module):

    def __init__(self, model):
      super().__init__()
      self.model = model

    @tf.function
    def f(self, imgs):
      return self.model(imgs, training=False, post_mode='global')

  imgs = tf.convert_to_tensor(imgs, dtype=tf.uint8)
  export_model = ExportModel(model)
  if FLAGS.saved_model_dir:
    tf.saved_model.save(
        export_model,
        FLAGS.saved_model_dir,
        signatures=export_model.f.get_concrete_function(
            tf.TensorSpec(shape=(None, None, None, 3), dtype=tf.uint8)))
    export_model = tf.saved_model.load(FLAGS.saved_model_dir)

  boxes, scores, classes, valid_len = export_model.f(imgs)

  # Visualize results.
  for i, img in enumerate(imgs):
    length = valid_len[i]
    img = inference.visualize_image(
        img,
        boxes[i].numpy()[:length],
        classes[i].numpy().astype(np.int)[:length],
        scores[i].numpy()[:length],
        label_map=config.label_map,
        min_score_thresh=config.nms_configs.score_thresh,
        max_boxes_to_draw=config.nms_configs.max_output_size)
    output_image_path = os.path.join(FLAGS.output_dir, str(i) + '.jpg')
    Image.fromarray(img).save(output_image_path)
    print('writing annotated image to %s' % output_image_path)
//  [python]
def init_tf(FLAGS):
    """
    Set global options for TensorFlow
    """
    gpus = tf.config.experimental.list_physical_devices('GPU')

    for gpu in gpus:
        tf.config.experimental.set_memory_growth(gpu, True)

    visible_gpus = []
    if gpus and not FLAGS.cpu:
        visible_gpus = gpus[hvd.local_rank()]
    tf.config.experimental.set_visible_devices(visible_gpus, 'GPU')

    if FLAGS.amp:
        policy = tf.keras.mixed_precision.Policy("mixed_float16")
        tf.keras.mixed_precision.set_global_policy(policy)

    tf.compat.v1.enable_eager_execution(FLAGS.run_eagerly)

    if FLAGS.tfdata_debug:
        tf.data.experimental.enable_debug_mode()

    if FLAGS.inter_op_parallelism:
        tf.config.threading.set_inter_op_parallelism_threads(FLAGS.inter_op_parallelism)

    if FLAGS.intra_op_parallelism:
        tf.config.threading.set_intra_op_parallelism_threads(FLAGS.intra_op_parallelism)

    tf.random.set_seed(hash((FLAGS.seed, hvd.rank())))
//  [python]
def init_tf(FLAGS):
    """
    Set global options for TensorFlow
    """
    gpus = tf.config.experimental.list_physical_devices('GPU')

    for gpu in gpus:
        tf.config.experimental.set_memory_growth(gpu, True)

    visible_gpus = []
    if gpus and not FLAGS.cpu:
        visible_gpus = gpus[hvd.local_rank()]
    tf.config.experimental.set_visible_devices(visible_gpus, 'GPU')

    if FLAGS.amp:
        policy = tf.keras.mixed_precision.Policy("mixed_float16")
        tf.keras.mixed_precision.set_global_policy(policy)

    tf.config.run_functions_eagerly(FLAGS.run_eagerly)

    if FLAGS.tfdata_debug:
        tf.data.experimental.enable_debug_mode()

    if FLAGS.inter_op_parallelism:
        tf.config.threading.set_inter_op_parallelism_threads(FLAGS.inter_op_parallelism)

    if FLAGS.intra_op_parallelism:
        tf.config.threading.set_intra_op_parallelism_threads(FLAGS.intra_op_parallelism)

    tf.random.set_seed(hash((FLAGS.seed, hvd.rank())))
//  [python]
def main(_):
  tf.compat.v1.disable_eager_execution()
  logging.set_verbosity(logging.INFO)
  tf.compat.v1.enable_v2_behavior()
  strategy = strategy_utils.get_strategy(FLAGS.tpu, FLAGS.use_gpu)

  gin.parse_config_files_and_bindings(FLAGS.gin_file, FLAGS.gin_bindings)
  logging.info('Gin bindings: %s', FLAGS.gin_bindings)
  logging.info('# Gin-Config:\n %s', gin.config.operative_config_str())

  train_eval(FLAGS.root_dir, strategy=strategy, reverb_port=FLAGS.reverb_port)
//  [python]
def main(_):
  tf.config.run_functions_eagerly(False)
  logging.set_verbosity(logging.INFO)
  tf.compat.v1.enable_v2_behavior()
  strategy = strategy_utils.get_strategy(FLAGS.tpu, FLAGS.use_gpu)

  gin.parse_config_files_and_bindings(FLAGS.gin_file, FLAGS.gin_bindings)
  logging.info('Gin bindings: %s', FLAGS.gin_bindings)
  logging.info('# Gin-Config:\n %s', gin.config.operative_config_str())

  train_eval(FLAGS.root_dir, strategy=strategy, reverb_port=FLAGS.reverb_port)
//  [python]
def configure_tensorflow(hparams):
    """Custom tensorflow import that configures proper flags, path settings,
    seeds, etc.

    Returns:
        module: Tensorflow as tf
    """
    set_tf_env_flags()
    tf = set_tf_expand_memory()
    tf.keras.backend.clear_session()
    tf.random.set_seed(hparams["seed"][0])
    tf.config.experimental_run_functions_eagerly(hparams["eager"][0])
    mixed_precision = set_mixed_precision() if hparams["mixed_precision"][0] else None

    return tf, mixed_precision
//  [python]
def configure_tensorflow(hparams):
    """Custom tensorflow import that configures proper flags, path settings,
    seeds, etc.

    Returns:
        module: Tensorflow as tf
    """
    set_tf_env_flags()
    tf = set_tf_expand_memory()
    tf.keras.backend.clear_session()
    tf.random.set_seed(hparams["seed"][0])
    tf.config.run_functions_eagerly(hparams["eager"][0])
    mixed_precision = set_mixed_precision() if hparams["mixed_precision"][0] else None

    return tf, mixed_precision
//  [python]
def build_tf2_model():
  """Build the tf2 model."""
  tf.compat.v1.enable_eager_execution(FLAGS.debug)
  if FLAGS.mixed_precision:
    # Use 'mixed_float16' if running on GPUs.
    policy = tf.keras.mixed_precision.Policy('mixed_float16')
    tf.keras.mixed_precision.set_global_policy(policy)

  model = effnetv2_model.get_model(
      FLAGS.model_name,
      FLAGS.hparam_str,
      include_top=True,
      weights=FLAGS.model_dir or 'imagenet')
  model.summary()
  return model
//  [python]
def build_tf2_model():
  """Build the tf2 model."""
  tf.config.run_functions_eagerly(FLAGS.debug)
  if FLAGS.mixed_precision:
    # Use 'mixed_float16' if running on GPUs.
    policy = tf.keras.mixed_precision.Policy('mixed_float16')
    tf.keras.mixed_precision.set_global_policy(policy)

  model = effnetv2_model.get_model(
      FLAGS.model_name,
      FLAGS.hparam_str,
      include_top=True,
      weights=FLAGS.model_dir or 'imagenet')
  model.summary()
  return model
//  [python]
def main(_):

  # pylint: disable=line-too-long
  # Prepare images and checkpoints: please run these commands in shell.
  # !mkdir tmp
  # !wget https://user-images.githubusercontent.com/11736571/77320690-099af300-6d37-11ea-9d86-24f14dc2d540.png -O tmp/img.png
  # !wget https://storage.googleapis.com/cloud-tpu-checkpoints/efficientdet/coco/efficientdet-d0.tar.gz -O tmp/efficientdet-d0.tar.gz
  # !tar zxf tmp/efficientdet-d0.tar.gz -C tmp
  imgs = [np.array(Image.open(FLAGS.image_path))]
  # Create model config.
  config = hparams_config.get_efficientdet_config(FLAGS.model_name)
  config.is_training_bn = False
  config.image_size = '1920x1280'
  config.nms_configs.score_thresh = 0.4
  config.nms_configs.max_output_size = 100
  config.override(FLAGS.hparams)

  # Use 'mixed_float16' if running on GPUs.
  policy = tf.keras.mixed_precision.Policy('float32')
  tf.keras.mixed_precision.set_global_policy(policy)
  tf.compat.v1.enable_eager_execution(FLAGS.debug)

  # Create and run the model.
  model = efficientdet_keras.EfficientDetModel(config=config)
  model.build((None, None, None, 3))
  model.load_weights(tf.train.latest_checkpoint(FLAGS.model_dir))
  model.summary(expand_nested=True)

  class ExportModel(tf.Module):

    def __init__(self, model):
      super().__init__()
      self.model = model

    @tf.function
    def f(self, imgs):
      return self.model(imgs, training=False, post_mode='global')

  imgs = tf.convert_to_tensor(imgs, dtype=tf.uint8)
  export_model = ExportModel(model)
  if FLAGS.saved_model_dir:
    tf.saved_model.save(
        export_model,
        FLAGS.saved_model_dir,
        signatures=export_model.f.get_concrete_function(
            tf.TensorSpec(shape=(None, None, None, 3), dtype=tf.uint8)))
    export_model = tf.saved_model.load(FLAGS.saved_model_dir)

  boxes, scores, classes, valid_len = export_model.f(imgs)

  # Visualize results.
  for i, img in enumerate(imgs):
    length = valid_len[i]
    img = inference.visualize_image(
        img,
        boxes[i].numpy()[:length],
        classes[i].numpy().astype(np.int)[:length],
        scores[i].numpy()[:length],
        label_map=config.label_map,
        min_score_thresh=config.nms_configs.score_thresh,
        max_boxes_to_draw=config.nms_configs.max_output_size)
    output_image_path = os.path.join(FLAGS.output_dir, str(i) + '.jpg')
    Image.fromarray(img).save(output_image_path)
    print('writing annotated image to %s' % output_image_path)
//  [python]
def main(_):

  # pylint: disable=line-too-long
  # Prepare images and checkpoints: please run these commands in shell.
  # !mkdir tmp
  # !wget https://user-images.githubusercontent.com/11736571/77320690-099af300-6d37-11ea-9d86-24f14dc2d540.png -O tmp/img.png
  # !wget https://storage.googleapis.com/cloud-tpu-checkpoints/efficientdet/coco/efficientdet-d0.tar.gz -O tmp/efficientdet-d0.tar.gz
  # !tar zxf tmp/efficientdet-d0.tar.gz -C tmp
  imgs = [np.array(Image.open(FLAGS.image_path))]
  # Create model config.
  config = hparams_config.get_efficientdet_config(FLAGS.model_name)
  config.is_training_bn = False
  config.image_size = '1920x1280'
  config.nms_configs.score_thresh = 0.4
  config.nms_configs.max_output_size = 100
  config.override(FLAGS.hparams)

  # Use 'mixed_float16' if running on GPUs.
  policy = tf.keras.mixed_precision.Policy('float32')
  tf.keras.mixed_precision.set_global_policy(policy)
  tf.config.run_functions_eagerly(FLAGS.debug)

  # Create and run the model.
  model = efficientdet_keras.EfficientDetModel(config=config)
  model.build((None, None, None, 3))
  model.load_weights(tf.train.latest_checkpoint(FLAGS.model_dir))
  model.summary(expand_nested=True)

  class ExportModel(tf.Module):

    def __init__(self, model):
      super().__init__()
      self.model = model

    @tf.function
    def f(self, imgs):
      return self.model(imgs, training=False, post_mode='global')

  imgs = tf.convert_to_tensor(imgs, dtype=tf.uint8)
  export_model = ExportModel(model)
  if FLAGS.saved_model_dir:
    tf.saved_model.save(
        export_model,
        FLAGS.saved_model_dir,
        signatures=export_model.f.get_concrete_function(
            tf.TensorSpec(shape=(None, None, None, 3), dtype=tf.uint8)))
    export_model = tf.saved_model.load(FLAGS.saved_model_dir)

  boxes, scores, classes, valid_len = export_model.f(imgs)

  # Visualize results.
  for i, img in enumerate(imgs):
    length = valid_len[i]
    img = inference.visualize_image(
        img,
        boxes[i].numpy()[:length],
        classes[i].numpy().astype(np.int)[:length],
        scores[i].numpy()[:length],
        label_map=config.label_map,
        min_score_thresh=config.nms_configs.score_thresh,
        max_boxes_to_draw=config.nms_configs.max_output_size)
    output_image_path = os.path.join(FLAGS.output_dir, str(i) + '.jpg')
    Image.fromarray(img).save(output_image_path)
    print('writing annotated image to %s' % output_image_path)
//  [python]
def testConcat(self):
    """Tests inconsistent behavior with `.concat()`.

    Bugs:   b/180921284
    Status: Missing error message
    Issue:  Running functions eagerly, calling `.concat` on a `tf.TensorArray`
            should raise an error but it does not.

    Error message:
      Expected error message is "Concatenating scalars in `tf.TensorArray` is
      unsupported in eager mode. Please use `.stack()` instead".

    Improve error message? Needed. (b/187851559)

    Notes:
    * Inconsistent behavior between eager and non-eager mode execution of the
      `tf.function` decorated function. In graph mode, the error is thrown.
    * We need to improve graph mode error message though. The error gets thrown
      is "Shapes must be equal rank, but are 1 and 0" and is hard to grasp.
    * Using `.stack()` as a workaround is working as intended:
      
      def f(x):
        return x.write(1, tf.constant([4, 5, 6]))

      ta = tf.TensorArray(dtype=tf.int32, dynamic_size=True, size=0)
      ta = ta.write(0, tf.constant([1, 2, 3]))
      f(ta).stack()  # <tf.Tensor: shape=(2, 3), dtype=int32,
                     # numpy=array([[1, 2, 3], [4, 5, 6]], dtype=int32)>
      
    """
    self.skipTest('b/180921284')
    try:
      tf.compat.v1.enable_eager_execution()

      @tf.function
      def f(x, y, z):
        ta = tf.TensorArray(size=3, dtype=tf.int32, element_shape=())
        ta = ta.write(0, x)
        ta = ta.write(1, y)
        ta = ta.write(2, z)
        return ta.concat()

      with self.assertRaisesWithPredicateMatch(
          BaseException,
          # TODO(hyey): Below is a placeholder error message of what we
          # probably want but it needs to be updated to specify what caused
          # the error and where.
          'Concatenating scalars in `tf.TensorArray` is unsupported in eager '
          'mode. Please use `.stack()` instead'):
        f(1, 2, 3)

    finally:
      tf.compat.v1.disable_eager_execution()
//  [python]
def testConcat(self):
    """Tests inconsistent behavior with `.concat()`.

    Bugs:   b/180921284
    Status: Missing error message
    Issue:  Running functions eagerly, calling `.concat` on a `tf.TensorArray`
            should raise an error but it does not.

    Error message:
      Expected error message is "Concatenating scalars in `tf.TensorArray` is
      unsupported in eager mode. Please use `.stack()` instead".

    Improve error message? Needed. (b/187851559)

    Notes:
    * Inconsistent behavior between eager and non-eager mode execution of the
      `tf.function` decorated function. In graph mode, the error is thrown.
    * We need to improve graph mode error message though. The error gets thrown
      is "Shapes must be equal rank, but are 1 and 0" and is hard to grasp.
    * Using `.stack()` as a workaround is working as intended:
      ```
      def f(x):
        return x.write(1, tf.constant([4, 5, 6]))

      ta = tf.TensorArray(dtype=tf.int32, dynamic_size=True, size=0)
      ta = ta.write(0, tf.constant([1, 2, 3]))
      f(ta).stack()  # <tf.Tensor: shape=(2, 3), dtype=int32,
                     # numpy=array([[1, 2, 3], [4, 5, 6]], dtype=int32)>
      ```
    """
    self.skipTest('b/180921284')
    try:
      tf.config.run_functions_eagerly(True)

      @tf.function
      def f(x, y, z):
        ta = tf.TensorArray(size=3, dtype=tf.int32, element_shape=())
        ta = ta.write(0, x)
        ta = ta.write(1, y)
        ta = ta.write(2, z)
        return ta.concat()

      with self.assertRaisesWithPredicateMatch(
          BaseException,
          # TODO(hyey): Below is a placeholder error message of what we
          # probably want but it needs to be updated to specify what caused
          # the error and where.
          'Concatenating scalars in `tf.TensorArray` is unsupported in eager '
          'mode. Please use `.stack()` instead'):
        f(1, 2, 3)

    finally:
      tf.config.run_functions_eagerly(False)
//  [python]
def build(self, params_override=None):
    """Build model and restore checkpoints."""
    params = copy.deepcopy(self.params)
    if params_override:
      params.update(params_override)
    config = hparams_config.get_efficientdet_config(self.model_name)
    config.override(params)
    if self.only_network:
      self.model = efficientdet_keras.EfficientDetNet(config=config)
    else:
      self.model = efficientdet_keras.EfficientDetModel(config=config)
    image_size = utils.parse_image_size(params['image_size'])
    self.model.build((self.batch_size, *image_size, 3))
    util_keras.restore_ckpt(self.model, self.ckpt_path,
                            self.params['moving_average_decay'],
                            skip_mismatch=False)
    if self.debug:
      tf.compat.v1.config.run_functions_eagerly(self.debug)
//  [python]
def build(self, params_override=None):
    """Build model and restore checkpoints."""
    params = copy.deepcopy(self.params)
    if params_override:
      params.update(params_override)
    config = hparams_config.get_efficientdet_config(self.model_name)
    config.override(params)
    if self.only_network:
      self.model = efficientdet_keras.EfficientDetNet(config=config)
    else:
      self.model = efficientdet_keras.EfficientDetModel(config=config)
    image_size = utils.parse_image_size(params['image_size'])
    self.model.build((self.batch_size, *image_size, 3))
    util_keras.restore_ckpt(self.model, self.ckpt_path,
                            self.params['moving_average_decay'],
                            skip_mismatch=False)
    if self.debug:
      tf.config.run_functions_eagerly(self.debug)
//  [python]
def argmin(in_tensor: Tensor, axis: Union[int, Iterable, tuple[int], None] = None):
    if BACKEND == BackEnd.NUMPY:
        return numpy.argmin(a=in_tensor, axis=axis)
    elif BACKEND == BackEnd.TENSORFLOW:
        return tf.argmin(input=in_tensor, axis=axis)
    else:
        raise NotImplementedError
//  [python]
def argmin(in_tensor: Tensor, axis: Union[int, Iterable, tuple[int], None] = None):
    if BACKEND == BackEnd.NUMPY:
        return numpy.argmin(a=in_tensor, axis=axis)
    elif BACKEND == BackEnd.TENSORFLOW:
        return tensorflow.math.argmin(input_tensor=in_tensor, axis=axis)
    else:
        raise NotImplementedError
//  [python]
def __init__(self, ckpt_path, debug, *args, **kwargs):
    """ Initialize the inference driver.

    Args:
      ckpt_path: checkpoint path, such as /tmp/efficientdet-d0/.
      debug: bool, if true, run in debug mode.
    """
    super().__init__(*args, **kwargs)
    params = copy.deepcopy(self.params)
    config = hparams_config.get_efficientdet_config(self.model_name)
    config.override(params)
    precision = utils.get_precision(config.strategy, config.mixed_precision)
    policy = tf.keras.mixed_precision.Policy(precision)
    tf.keras.mixed_precision.set_global_policy(policy)
    self.model = efficientdet_keras.EfficientDetModel(config=config)
    image_size = utils.parse_image_size(config.image_size)
    self.model.build((self.batch_size, *image_size, 3))
    util_keras.restore_ckpt(
        self.model, ckpt_path, config.moving_average_decay, skip_mismatch=False)
    self.debug = debug
    if debug:
      tf.config.run_functions_eagerly(True)
//  [python]
def __init__(self, ckpt_path, debug, *args, **kwargs):
    """ Initialize the inference driver.

    Args:
      ckpt_path: checkpoint path, such as /tmp/efficientdet-d0/.
      debug: bool, if true, run in debug mode.
    """
    super().__init__(*args, **kwargs)
    params = copy.deepcopy(self.params)
    config = hparams_config.get_efficientdet_config(self.model_name)
    config.override(params)
    precision = utils.get_precision(config.strategy, config.mixed_precision)
    policy = tf.keras.mixed_precision.Policy(precision)
    tf.keras.mixed_precision.set_global_policy(policy)
    self.model = efficientdet_keras.EfficientDetModel(config=config)
    image_size = utils.parse_image_size(config.image_size)
    self.model.build((self.batch_size, *image_size, 3))
    util_keras.restore_ckpt(
        self.model, ckpt_path, config.moving_average_decay, skip_mismatch=False)
    self.debug = debug
    if debug:
      tf.config.run_functions_eagerly(debug)
//  [python]
def avg_iou_total(self, boxes, clusters):
    clusters = tf.convert_to_tensor(clusters)
    dists = 1 - self.iou(boxes, clusters)
    assignments = tf.argmin(dists, axis=-1)
    ious = self._avg_iou(boxes, clusters, assignments)
    return clusters, assignments, ious
//  [python]
def avg_iou_total(self, boxes, clusters):
    clusters = tf.convert_to_tensor(clusters)
    dists = 1 - self.iou(boxes, clusters)
    assignments = tf.math.argmin(dists, axis=-1)
    ious = self._avg_iou(boxes, clusters, assignments)
    return clusters, assignments, ious
//  [python]
def fixture_device_api_main_func():
    """Test function generator which generates C Device API calls"""

    # Ideally we should have a sample Target registered here
    # but we're going to re-use this for now
    pytest.importorskip("ethosu.vela")

    # pylint: disable=import-outside-toplevel
    import tensorflow as tf
    import tflite.Model
    from tvm.relay.op.contrib.ethosu import partition_for_ethosu

    from tests.python.contrib.test_ethosu.infra import (
        create_test_runner,
        generate_ref_data_tflite,
    )

    # pylint: enable=import-outside-toplevel

    tf.compat.v1.disable_eager_execution()

    class Model(tf.Module):
        @tf.function
        def tf_function(self, x):
            return tf.nn.max_pool(x, [1, 2], [1, 2], "SAME")

    def representative_dataset():
        for _ in range(100):
            data = np.random.rand(1, 3, 4, 3)
            yield [data.astype(np.float32)]

    model = Model()
    concrete_func = model.tf_function.get_concrete_function(
        tf.TensorSpec([1, 3, 4, 3], dtype=tf.float32)
    )

    converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])
    converter.optimizations = [tf.lite.Optimize.DEFAULT]
    converter.representative_dataset = representative_dataset
    converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
    converter.inference_input_type = tf.int8
    converter.inference_output_type = tf.int8

    tflite_graph = converter.convert()
    tflite_model = tflite.Model.Model.GetRootAsModel(tflite_graph, 0)

    relay_module, params = relay.frontend.from_tflite(
        tflite_model,
        shape_dict={"x": [1, 3, 4, 3]},
        dtype_dict={"x": "int8"},
    )
    mod = partition_for_ethosu(relay_module, params)

    # Generate reference data
    input_data, output_data = generate_ref_data_tflite(tflite_graph)

    def compile_to_main_func(interface_api="c", use_unpacked_api=True):
        test_runner = create_test_runner()
        compiled_models = compile_models(
            models=AOTTestModel(
                module=mod,
                inputs=input_data,
                outputs=output_data,
            ),
            interface_api=interface_api,
            use_unpacked_api=use_unpacked_api,
            workspace_byte_alignment=16,
            pass_config=test_runner.pass_config,
        )
        main_ir_module = compiled_models[0].executor_factory.lowered_ir_mods.items()[0][1]
        main_func = main_ir_module["__tvm_main__"]
        return main_func

    return compile_to_main_func