Fix lite backbone builder.

Fix #357
Also add a few simple code cleanup.

Author: mingxingtan

efficientdet/backbone/efficientnet_builder_test.py

@@ -38,7 +38,7 @@ def _test_model_params(self,
         images,
         model_name=model_name,
         override_params=override_params,
-        training=True,
+        training=False,
         features_only=features_only,
         pooled_features_only=pooled_features_only)
     num_params = np.sum([np.prod(v.shape) for v in tf.trainable_variables()])
@@ -91,15 +91,15 @@ def test_efficientnet_b0_fails_if_both_features_requested(self):
       efficientnet_builder.build_model(
           None,
           model_name='efficientnet-b0',
-          training=True,
+          training=False,
           features_only=True,
           pooled_features_only=True)
 
   def test_efficientnet_b0_base(self):
     # Creates a base model using the model configuration.
     images = tf.zeros((1, 224, 224, 3), dtype=tf.float32)
     _, endpoints = efficientnet_builder.build_model_base(
-        images, model_name='efficientnet-b0', training=True)
+        images, model_name='efficientnet-b0', training=False)
 
     # reduction_1 to reduction_5 should be in endpoints
     self.assertIn('reduction_1', endpoints)

efficientdet/backbone/efficientnet_lite_builder.py

@@ -161,13 +161,12 @@ def build_model(images,
         f.write('global_params= %s\n\n' % str(global_params))
         f.write('blocks_args= %s\n\n' % str(blocks_args))
 
-  with tf.variable_scope(model_name):
-    model = efficientnet_model.Model(blocks_args, global_params)
-    outputs = model(
-        images,
-        training=training,
-        features_only=features_only,
-        pooled_features_only=pooled_features_only)
+  model = efficientnet_model.Model(blocks_args, global_params, model_name)
+  outputs = model(
+      images,
+      training=training,
+      features_only=features_only,
+      pooled_features_only=pooled_features_only)
   if features_only:
     outputs = tf.identity(outputs, 'features')
   elif pooled_features_only:
@@ -202,9 +201,8 @@ def build_model_base(images, model_name, training, override_params=None):
 
   blocks_args, global_params = get_model_params(model_name, override_params)
 
-  with tf.variable_scope(model_name):
-    model = efficientnet_model.Model(blocks_args, global_params)
-    features = model(images, training=training, features_only=True)
+  model = efficientnet_model.Model(blocks_args, global_params, model_name)
+  features = model(images, training=training, features_only=True)
 
   features = tf.identity(features, 'features')
   return features, model.endpoints

efficientdet/backbone/efficientnet_lite_builder_test.py

@@ -38,7 +38,7 @@ def _test_model_params(self,
         images,
         model_name=model_name,
         override_params=override_params,
-        training=True,
+        training=False,
         features_only=features_only,
         pooled_features_only=pooled_features_only)
     num_params = np.sum([np.prod(v.shape) for v in tf.trainable_variables()])

efficientdet/backbone/efficientnet_model_test.py

@@ -35,7 +35,7 @@ def test_bottleneck_block(self):
         0,
         'channels_last',
         num_classes=10,
-        batch_norm=utils.TpuBatchNormalization)
+        batch_norm=utils.batch_norm_class(False))
     blocks_args = [
         efficientnet_model.BlockArgs(
             kernel_size=3,

efficientdet/utils.py

@@ -236,7 +236,7 @@ def call(self, *args, **kwargs):
     return outputs
 
 
-def batch_norm_class(is_training, use_tpu=False,):
+def batch_norm_class(is_training, use_tpu=False):
   if is_training and use_tpu:
     return TpuBatchNormalization
   else: