ready for review

official/vision/beta/projects/yolo/modeling/backbones/darknet.py

@@ -77,11 +77,11 @@ class layer_factory(object):
   """
   def __init__(self):
     self._layer_dict = {
-        "DarkConv": (nn_blocks.DarkConv, self.darkconv_config_todict),
+        "ConvBN": (nn_blocks.ConvBN, self.ConvBN_config_todict),
         "MaxPool": (tf.keras.layers.MaxPool2D, self.maxpool_config_todict)
     }
 
-  def darkconv_config_todict(self, config, kwargs):
+  def ConvBN_config_todict(self, config, kwargs):
     dictvals = {
         "filters": config.filters,
         "kernel_size": config.kernel_size,
@@ -124,7 +124,7 @@ CSPDARKNET53 = {
     "splits": {"backbone_split": 106,
                "neck_split": 138},
     "backbone": [
-        ["DarkConv", None, 1, False, 32, None, 3, 1, "same", "mish", -1, 0, False],
+        ["ConvBN", None, 1, False, 32, None, 3, 1, "same", "mish", -1, 0, False],
         ["DarkRes", "csp", 1, True, 64, None, None, None, None, "mish", -1, 1, False],
         ["DarkRes", "csp", 2, False, 128, None, None, None, None, "mish", -1, 2, False],
         ["DarkRes", "csp", 8, False, 256, None, None, None, None, "mish", -1, 3, True],
@@ -137,7 +137,7 @@ DARKNET53 = {
     "list_names": LISTNAMES,
     "splits": {"backbone_split": 76},
     "backbone": [
-        ["DarkConv", None, 1, False, 32, None, 3, 1, "same", "leaky", -1, 0, False],
+        ["ConvBN", None, 1, False, 32, None, 3, 1, "same", "leaky", -1, 0, False],
         ["DarkRes", "residual", 1, True, 64, None, None, None, None, "leaky", -1, 1, False],
         ["DarkRes", "residual", 2, False, 128, None, None, None, None, "leaky", -1, 2, False],
         ["DarkRes", "residual", 8, False, 256, None, None, None, None, "leaky", -1, 3, True],
@@ -150,12 +150,12 @@ CSPDARKNETTINY = {
     "list_names": LISTNAMES,
     "splits": {"backbone_split": 28},
     "backbone": [
-        ["DarkConv", None, 1, False, 32, None, 3, 2, "same", "leaky", -1, 0, False],
-        ["DarkConv", None, 1, False, 64, None, 3, 2, "same", "leaky", -1, 1, False],
+        ["ConvBN", None, 1, False, 32, None, 3, 2, "same", "leaky", -1, 0, False],
+        ["ConvBN", None, 1, False, 64, None, 3, 2, "same", "leaky", -1, 1, False],
         ["CSPTiny", "csp_tiny", 1, False, 64, None, 3, 2, "same", "leaky", -1, 2, False],
         ["CSPTiny", "csp_tiny", 1, False, 128, None, 3, 2, "same", "leaky", -1, 3, False],
         ["CSPTiny", "csp_tiny", 1, False, 256, None, 3, 2, "same", "leaky", -1, 4, True],
-        ["DarkConv", None, 1, False, 512, None, 3, 1, "same", "leaky", -1, 5, True],
+        ["ConvBN", None, 1, False, 512, None, 3, 1, "same", "leaky", -1, 5, True],
     ]
 }
 
@@ -163,7 +163,7 @@ DARKNETTINY = {
     "list_names": LISTNAMES,
     "splits": {"backbone_split": 14},
     "backbone": [
-        ["DarkConv", None, 1, False, 16, None, 3, 1, "same", "leaky", -1, 0, False],
+        ["ConvBN", None, 1, False, 16, None, 3, 1, "same", "leaky", -1, 0, False],
         ["DarkTiny", "tiny", 1, True, 32, None, 3, 2, "same", "leaky", -1, 1, False],
         ["DarkTiny", "tiny", 1, True, 64, None, 3, 2, "same", "leaky", -1, 2, False],
         ["DarkTiny", "tiny", 1, False, 128, None, 3, 2, "same", "leaky", -1, 3, False],
@@ -292,27 +292,28 @@ class Darknet(ks.Model):
 
   def _csp_stack(self, inputs, config, name):
     if config.bottleneck:
-      csp_filter_reduce = 1
-      residual_filter_reduce = 2
+      csp_filter_scale = 1
+      residual_filter_scale = 2
       scale_filters = 1
     else:
-      csp_filter_reduce = 2
-      residual_filter_reduce = 1
+      csp_filter_scale = 2
+      residual_filter_scale = 1
       scale_filters = 2
     self._default_dict["activation"] = self._get_activation(config.activation)
     self._default_dict["name"] = f"{name}_csp_down"
-    x, x_route = nn_blocks.CSPDownSample(filters=config.filters,
-                                         filter_reduce=csp_filter_reduce,
-                                         **self._default_dict)(inputs)
+    x, x_route = nn_blocks.CSPRoute(filters=config.filters,
+                                    filter_scale=csp_filter_scale,
+                                    downsample=True,
+                                    **self._default_dict)(inputs)
     for i in range(config.repetitions):
       self._default_dict["name"] = f"{name}_{i}"
       x = nn_blocks.DarkResidual(filters=config.filters // scale_filters,
-                                 filter_scale=residual_filter_reduce,
+                                 filter_scale=residual_filter_scale,
                                  **self._default_dict)(x)
 
     self._default_dict["name"] = f"{name}_csp_connect"
     output = nn_blocks.CSPConnect(filters=config.filters,
-                                  filter_reduce=csp_filter_reduce,
+                                  filter_scale=csp_filter_scale,
                                   **self._default_dict)([x, x_route])
     self._default_dict["activation"] = self._activation
     self._default_dict["name"] = None
@@ -335,7 +336,7 @@ class Darknet(ks.Model):
                                   name=f"{name}_tiny/pool")(inputs)
     self._default_dict["activation"] = self._get_activation(config.activation)
     self._default_dict["name"] = f"{name}_tiny/conv"
-    x = nn_blocks.DarkConv(filters=config.filters,
+    x = nn_blocks.ConvBN(filters=config.filters,
                            kernel_size=(3, 3),
                            strides=(1, 1),
                            padding='same',

official/vision/beta/projects/yolo/modeling/layers/nn_blocks_test.py

@@ -7,14 +7,14 @@ from official.vision.beta.projects.yolo.modeling.layers import nn_blocks
 
 
 
-class CSPConnect(tf.test.TestCase, parameterized.TestCase):
+class CSPConnectTest(tf.test.TestCase, parameterized.TestCase):
 
   @parameterized.named_parameters(("same", 224, 224, 64, 1),
                                   ("downsample", 224, 224, 64, 2))
   def test_pass_through(self, width, height, filters, mod):
     x = ks.Input(shape=(width, height, filters))
-    test_layer = nn_blocks.CSPDownSample(filters=filters, filter_reduce=mod)
-    test_layer2 = nn_blocks.CSPConnect(filters=filters, filter_reduce=mod)
+    test_layer = nn_blocks.CSPRoute(filters=filters, filter_scale=mod)
+    test_layer2 = nn_blocks.CSPConnect(filters=filters, filter_scale=mod)
     outx, px = test_layer(x)
     outx = test_layer2([outx, px])
     print(outx)
@@ -29,8 +29,8 @@ class CSPConnect(tf.test.TestCase, parameterized.TestCase):
   def test_gradient_pass_though(self, filters, width, height, mod):
     loss = ks.losses.MeanSquaredError()
     optimizer = ks.optimizers.SGD()
-    test_layer = nn_blocks.CSPDownSample(filters, filter_reduce=mod)
-    path_layer = nn_blocks.CSPConnect(filters, filter_reduce=mod)
+    test_layer = nn_blocks.CSPRoute(filters, filter_scale=mod)
+    path_layer = nn_blocks.CSPConnect(filters, filter_scale=mod)
 
     init = tf.random_normal_initializer()
     x = tf.Variable(
@@ -49,13 +49,13 @@ class CSPConnect(tf.test.TestCase, parameterized.TestCase):
 
     self.assertNotIn(None, grad)
 
-class CSPDownSample(tf.test.TestCase, parameterized.TestCase):
+class CSPRouteTest(tf.test.TestCase, parameterized.TestCase):
 
   @parameterized.named_parameters(("same", 224, 224, 64, 1),
                                   ("downsample", 224, 224, 64, 2))
   def test_pass_through(self, width, height, filters, mod):
     x = ks.Input(shape=(width, height, filters))
-    test_layer = nn_blocks.CSPDownSample(filters=filters, filter_reduce=mod)
+    test_layer = nn_blocks.CSPRoute(filters=filters, filter_scale=mod)
     outx, px = test_layer(x)
     print(outx)
     print(outx.shape.as_list())
@@ -69,8 +69,8 @@ class CSPDownSample(tf.test.TestCase, parameterized.TestCase):
   def test_gradient_pass_though(self, filters, width, height, mod):
     loss = ks.losses.MeanSquaredError()
     optimizer = ks.optimizers.SGD()
-    test_layer = nn_blocks.CSPDownSample(filters, filter_reduce=mod)
-    path_layer = nn_blocks.CSPConnect(filters, filter_reduce=mod)
+    test_layer = nn_blocks.CSPRoute(filters, filter_scale=mod)
+    path_layer = nn_blocks.CSPConnect(filters, filter_scale=mod)
 
     init = tf.random_normal_initializer()
     x = tf.Variable(
@@ -89,7 +89,75 @@ class CSPDownSample(tf.test.TestCase, parameterized.TestCase):
 
     self.assertNotIn(None, grad)
 
-class DarkConvTest(tf.test.TestCase, parameterized.TestCase):
+class CSPStackTest(tf.test.TestCase, parameterized.TestCase):
+  def build_layer(self, layer_type, filters, filter_scale, count, stack_type, downsample):
+    if stack_type != None:
+      layers = []
+      if layer_type == "residual":
+        for _ in range(count): 
+          layers.append(nn_blocks.DarkResidual(filters = filters // filter_scale, filter_scale = filter_scale))
+      else:
+        for _ in range(count): 
+          layers.append(nn_blocks.ConvBN(filters = filters))
+
+      if stack_type == "model":
+        layers = tf.keras.Sequential(layers=layers)
+    else:
+      layers = None
+    
+    stack = nn_blocks.CSPStack(filters = filters,
+                           filter_scale = filter_scale, 
+                           downsample = downsample,
+                           model_to_wrap = layers)
+    return stack
+
+  @parameterized.named_parameters(("no_stack", 224, 224, 64, 2, "residual", None, 0, True),
+                                  ("residual_stack", 224, 224, 64, 2, "residual", "list", 2, True), 
+                                  ("conv_stack", 224, 224, 64, 2, "conv", "list", 3, False), 
+                                  ("callable_no_scale", 224, 224, 64, 1, "residual", "model", 5, False))
+  def test_pass_through(self, width, height, filters, mod, layer_type, stack_type, count, downsample):
+    x = ks.Input(shape=(width, height, filters))
+    test_layer = self.build_layer(layer_type, filters, mod, count, stack_type, downsample)
+    outx = test_layer(x)
+    print(outx)
+    print(outx.shape.as_list())
+    if downsample:
+      self.assertAllEqual(
+          outx.shape.as_list(),
+          [None, width//2, height//2, filters])
+    else:
+      self.assertAllEqual(
+          outx.shape.as_list(),
+          [None, width, height, filters])
+
+
+  @parameterized.named_parameters(("no_stack", 224, 224, 64, 2, "residual", None, 0, True),
+                                  ("residual_stack", 224, 224, 64, 2, "residual", "list", 2, True), 
+                                  ("conv_stack", 224, 224, 64, 2, "conv", "list", 3, False), 
+                                  ("callable_no_scale", 224, 224, 64, 1, "residual", "model", 5, False))
+  def test_gradient_pass_though(self, width, height, filters, mod, layer_type, stack_type, count, downsample):
+    loss = ks.losses.MeanSquaredError()
+    optimizer = ks.optimizers.SGD()
+
+    init = tf.random_normal_initializer()
+    x = tf.Variable(initial_value=init(shape=(1, width, height, filters), dtype=tf.float32))
+
+    if not downsample:
+      y = tf.Variable(initial_value=init(shape=(1, width, height, filters), dtype=tf.float32))
+    else:
+      y = tf.Variable(initial_value=init(shape=(1, width//2, height//2, filters), dtype=tf.float32))
+    test_layer = self.build_layer(layer_type, filters, mod, count, stack_type, downsample)
+
+    with tf.GradientTape() as tape:
+      x_hat = test_layer(x)
+      grad_loss = loss(x_hat, y)
+    grad = tape.gradient(grad_loss, test_layer.trainable_variables)
+    optimizer.apply_gradients(zip(grad, test_layer.trainable_variables))
+
+    self.assertNotIn(None, grad)
+
+
+class ConvBNTest(tf.test.TestCase, parameterized.TestCase):
 
   @parameterized.named_parameters(
       ("valid", (3, 3), "valid", (1, 1)), ("same", (3, 3), "same", (1, 1)),
@@ -100,7 +168,7 @@ class DarkConvTest(tf.test.TestCase, parameterized.TestCase):
     else:
       pad_const = 0
     x = ks.Input(shape=(224, 224, 3))
-    test_layer = nn_blocks.DarkConv(filters=64,
+    test_layer = nn_blocks.ConvBN(filters=64,
                           kernel_size=kernel_size,
                           padding=padding,
                           strides=strides,
@@ -120,7 +188,7 @@ class DarkConvTest(tf.test.TestCase, parameterized.TestCase):
     loss = ks.losses.MeanSquaredError()
     optimizer = ks.optimizers.SGD()
     with tf.device("/CPU:0"):
-      test_layer = nn_blocks.DarkConv(filters, kernel_size=(3, 3), padding="same")
+      test_layer = nn_blocks.ConvBN(filters, kernel_size=(3, 3), padding="same")
 
     init = tf.random_normal_initializer()
     x = tf.Variable(initial_value=init(shape=(1, 224, 224,