Add option to set channels_per_group_list on ResNet.

PiperOrigin-RevId: 274159409
Change-Id: If31b33c236c5153aa1ad65b87293dfbca44674a7

sonnet/src/nets/resnet.py

@@ -233,6 +233,7 @@ def __init__(self,
                num_classes: int,
                bn_config: Optional[Mapping[Text, float]] = None,
                resnet_v2: bool = False,
+               channels_per_group_list: Sequence[int] = (256, 512, 1024, 2048),
                name: Optional[Text] = None):
     """Constructs a ResNet model.
 
@@ -245,6 +246,8 @@ def __init__(self,
         `0.9` and `eps` is `1e-5`.
       resnet_v2: Whether to use the v1 or v2 ResNet implementation. Defaults to
         False.
+      channels_per_group_list: A sequence of length 4 that indicates the number
+        of channels used for each block in each group.
       name: Name of the module.
     """
     super(ResNet, self).__init__(name=name)
@@ -264,6 +267,13 @@ def __init__(self,
               len(blocks_per_group_list)))
     self._blocks_per_group_list = blocks_per_group_list
 
+    # Number of channels in each group for ResNet.
+    if len(channels_per_group_list) != 4:
+      raise ValueError(
+          "`channels_per_group_list` must be of length 4 not {}".format(
+              len(channels_per_group_list)))
+    self._channels_per_group_list = channels_per_group_list
+
     self._initial_conv = conv.Conv2D(
         output_channels=64,
         kernel_shape=7,
@@ -283,7 +293,7 @@ def __init__(self,
     for i in range(4):
       self._block_groups.append(
           BlockGroup(
-              channels=256 * 2**i,
+              channels=self._channels_per_group_list[i],
               num_blocks=self._blocks_per_group_list[i],
               stride=strides[i],
               bn_config=bn_config,

sonnet/src/nets/resnet_test.py

@@ -55,13 +55,23 @@ def test_tf_function(self, resnet_v2):
     self.assertAllEqual(model(image, is_training=True).numpy(), logits.numpy())
 
   @parameterized.parameters(3, 5)
-  def test_error_incorrect_args(self, list_length):
+  def test_error_incorrect_args_block_list(self, list_length):
     block_list = [i for i in range(list_length)]
     with self.assertRaisesRegexp(
         ValueError, "blocks_per_group_list` must be of length 4 not {}".format(
             list_length)):
       resnet.ResNet(block_list, 10, {"decay_rate": 0.9, "eps": 1e-5})
 
+  @parameterized.parameters(3, 5)
+  def test_error_incorrect_args_channel_list(self, list_length):
+    channel_list = [i for i in range(list_length)]
+    with self.assertRaisesRegexp(
+        ValueError,
+        "channels_per_group_list` must be of length 4 not {}".format(
+            list_length)):
+      resnet.ResNet([1, 1, 1, 1], 10, {"decay_rate": 0.9, "eps": 1e-5},
+                    channels_per_group_list=channel_list)
+
   def test_v2_throws(self):
     resnet.TESTONLY_ENABLE_RESNET_V2 = False
     with self.assertRaisesRegexp(NotImplementedError, "please use v1"):