Repository: tensorflow/models
File: research/delf/delf/python/training/model/delg_model_test.py (Python)

Code before:
from absl.testing import parameterized
import tensorflow as tf

from delf.python.training.model import delg_model


class DelgTest(tf.test.TestCase, parameterized.TestCase):

  @parameterized.named_parameters(
      ('block3_stridesTrue', True),
      ('block3_stridesFalse', False),
  )
  def test_forward_pass(self, block3_strides):
    image_size = 321
    num_classes = 1000
    batch_size = 2
    input_shape = (batch_size, image_size, image_size, 3)
    local_feature_dim = 64
    feature_map_size = image_size // 16 # reduction factor for resnet50.
    if block3_strides:
        feature_map_size //= 2

    model = delg_model.Delg(block3_strides=block3_strides,
                            use_dim_reduction=True,
                            reduced_dimension=local_feature_dim
                            )
    model.init_classifiers(num_classes)

    images = tf.random.uniform(input_shape, minval=-1.0, maxval=1.0, seed=0)

    # Run a complete forward pass of the model.
    global_feature, attn_scores, local_features = model.build_call(images)

    self.assertAllEqual(global_feature.shape, (batch_size, 2048))
    self.assertAllEqual(
        attn_scores.shape,
        (batch_size, feature_map_size, feature_map_size, 1)
    )
    self.assertAllEqual(
        local_features.shape,
        (batch_size, feature_map_size, feature_map_size, local_feature_dim)
    )

  @parameterized.named_parameters(
      ('block3_stridesTrue', True),
      ('block3_stridesFalse', False),
  )
  def test_build_model(self, block3_strides):
    image_size = 321
    num_classes = 1000
    batch_size = 2

Review comment (refactor):
minor formatting nit: can you please move this close-parentheses to the previous line? Usually, per convention we leave the close-parentheses on the next line only if the previous one ends with a comma -- eg, replacing `=local_feature_dim` by `=local_feature_dim,` would also be fine. But maybe to keep conventions in these files, let's make it `=local_feature_dim)`

Code after:
import tensorflow as tf

from delf.python.training.model import delg_model


class DelgTest(tf.test.TestCase, parameterized.TestCase):

  @parameterized.named_parameters(
      ('block3_stridesTrue', True),
      ('block3_stridesFalse', False),
  )
  def test_forward_pass(self, block3_strides):
    image_size = 321
    num_classes = 1000
    batch_size = 2
    input_shape = (batch_size, image_size, image_size, 3)
    local_feature_dim = 64
    feature_map_size = image_size // 16 # reduction factor for resnet50.
    if block3_strides:
      feature_map_size //= 2

    model = delg_model.Delg(block3_strides=block3_strides,
                            use_dim_reduction=True,
                            reduced_dimension=local_feature_dim)
    model.init_classifiers(num_classes)

    images = tf.random.uniform(input_shape, minval=-1.0, maxval=1.0, seed=0)

    # Run a complete forward pass of the model.
    global_feature, attn_scores, local_features = model.build_call(images)

    self.assertAllEqual(global_feature.shape, (batch_size, 2048))
    self.assertAllEqual(
        attn_scores.shape,
        (batch_size, feature_map_size, feature_map_size, 1))
    self.assertAllEqual(
        local_features.shape,
        (batch_size, feature_map_size, feature_map_size, local_feature_dim))

  @parameterized.named_parameters(
      ('block3_stridesTrue', True),
      ('block3_stridesFalse', False),
  )
  def test_build_model(self, block3_strides):
    image_size = 321
    num_classes = 1000
    batch_size = 2
    input_shape = (batch_size, image_size, image_size, 3)

    model = delg_model.Delg(
        block3_strides=block3_strides,