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

Code before:
"""Tests for the DELG model."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from math import floor

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 = floor(image_size/16) # reduction factor for resnet50
    if block3_strides:
        feature_map_size //= 2

    model = delg_model.Delg(block3_strides=block3_strides,
                            name='DELG',
                            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)
    labels = tf.random.uniform((batch_size,),
                               minval=0,
                               maxval=model.num_classes - 1,
                               dtype=tf.int64)

    # 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),

Review comment (suggestion):
instead of `floor`, I think we could simply use `int`? Then, we no longer need the new import for `floor`

Code after:
"""Tests for the DELG model."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

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):