Internal change

PiperOrigin-RevId: 453477038

official/vision/modeling/backbones/spinenet.py

@@ -199,15 +199,11 @@ class SpineNet(tf.keras.Model):
     self._use_sync_bn = use_sync_bn
     self._norm_momentum = norm_momentum
     self._norm_epsilon = norm_epsilon
-    if activation == 'relu':
-      self._activation_fn = tf.nn.relu
-    elif activation == 'swish':
-      self._activation_fn = tf.nn.swish
-    else:
-      raise ValueError('Activation {} not implemented.'.format(activation))
     self._init_block_fn = 'bottleneck'
     self._num_init_blocks = 2
 
+    self._set_activation_fn(activation)
+
     if use_sync_bn:
       self._norm = layers.experimental.SyncBatchNormalization
     else:
@@ -232,6 +228,14 @@ class SpineNet(tf.keras.Model):
     self._output_specs = {l: endpoints[l].get_shape() for l in endpoints}
     super(SpineNet, self).__init__(inputs=inputs, outputs=endpoints)
 
+  def _set_activation_fn(self, activation):
+    if activation == 'relu':
+      self._activation_fn = tf.nn.relu
+    elif activation == 'swish':
+      self._activation_fn = tf.nn.swish
+    else:
+      raise ValueError('Activation {} not implemented.'.format(activation))
+
   def _block_group(self,
                    inputs: tf.Tensor,
                    filters: int,

official/vision/modeling/backbones/spinenet_test.py

@@ -122,6 +122,18 @@ class SpineNetTest(parameterized.TestCase, tf.test.TestCase):
     # If the serialization was successful, the new config should match the old.
     self.assertAllEqual(network.get_config(), new_network.get_config())
 
+  @parameterized.parameters(
+      ('relu', tf.nn.relu),
+      ('swish', tf.nn.swish)
+  )
+  def test_activation(self, activation, activation_fn):
+    model = spinenet.SpineNet(activation=activation)
+    self.assertEqual(model._activation_fn, activation_fn)
+
+  def test_invalid_activation_raises_valurerror(self):
+    with self.assertRaises(ValueError):
+      spinenet.SpineNet(activation='invalid_activation_name')
+
 
 if __name__ == '__main__':
   tf.test.main()