Refactored ResNet code and added additional architectures. (#6316)

* Refactored ResNet code and added additional architectures.

* Added numerical layer names instead of alphabetical.

* Change dash to underscore.

* Corrected return statement.

* Use conv_strides argument.

* Set classes=10

* Use partial to reduce code duplication.

official/resnet/keras/resnet_cifar_model.py

@@ -23,6 +23,7 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import functools
 import tensorflow as tf
 from tensorflow.python.keras import backend
 from tensorflow.python.keras import layers
@@ -47,7 +48,7 @@ def identity_building_block(input_tensor,
         middle conv layer at main path
     filters: list of integers, the filters of 3 conv layer at main path
     stage: integer, current stage label, used for generating layer names
-    block: 'a','b'..., current block label, used for generating layer names
+    block: current block label, used for generating layer names
     training: Only used if training keras model with Estimator.  In other
       scenarios it is handled automatically.
 
@@ -111,7 +112,7 @@ def conv_building_block(input_tensor,
         middle conv layer at main path
     filters: list of integers, the filters of 3 conv layer at main path
     stage: integer, current stage label, used for generating layer names
-    block: 'a','b'..., current block label, used for generating layer names
+    block: current block label, used for generating layer names
     strides: Strides for the first conv layer in the block.
     training: Only used if training keras model with Estimator.  In other
       scenarios it is handled automatically.
@@ -176,10 +177,49 @@ def conv_building_block(input_tensor,
   return x
 
 
-def resnet56(classes=100, training=None):
-  """Instantiates the ResNet56 architecture.
+def resnet_block(input_tensor,
+                 size,
+                 kernel_size,
+                 filters,
+                 stage,
+                 conv_strides=(2, 2),
+                 training=None):
+  """A block which applies conv followed by multiple identity blocks.
 
   Arguments:
+    input_tensor: input tensor
+    size: integer, number of constituent conv/identity building blocks.
+    A conv block is applied once, followed by (size - 1) identity blocks.
+    kernel_size: default 3, the kernel size of
+        middle conv layer at main path
+    filters: list of integers, the filters of 3 conv layer at main path
+    stage: integer, current stage label, used for generating layer names
+    conv_strides: Strides for the first conv layer in the block.
+    training: Only used if training keras model with Estimator.  In other
+      scenarios it is handled automatically.
+
+  Returns:
+    Output tensor after applying conv and identity blocks.
+  """
+
+  x = conv_building_block(input_tensor, kernel_size, filters, stage=stage,
+                          strides=conv_strides, block='block_0',
+                          training=training)
+  for i in range(size - 1):
+    x = identity_building_block(x, kernel_size, filters, stage=stage,
+                                block='block_%d' % (i + 1), training=training)
+  return x
+
+def resnet(num_blocks, classes=10, training=None):
+  """Instantiates the ResNet architecture.
+
+  Arguments:
+    num_blocks: integer, the number of conv/identity blocks in each block.
+      The ResNet contains 3 blocks with each block containing one conv block
+      followed by (layers_per_block - 1) number of idenity blocks. Each
+      conv/idenity block has 2 convolutional layers. With the input
+      convolutional layer and the pooling layer towards the end, this brings
+      the total size of the network to (6*num_blocks + 2)
     classes: optional number of classes to classify images into
     training: Only used if training keras model with Estimator.  In other
     scenarios it is handled automatically.
@@ -187,6 +227,7 @@ def resnet56(classes=100, training=None):
   Returns:
     A Keras model instance.
   """
+
   input_shape = (32, 32, 3)
   img_input = layers.Input(shape=input_shape)
 
@@ -214,62 +255,14 @@ def resnet56(classes=100, training=None):
                                              x, training=training)
   x = tf.keras.layers.Activation('relu')(x)
 
-  x = conv_building_block(x, 3, [16, 16], stage=2, block='a', strides=(1, 1),
-                          training=training)
-  x = identity_building_block(x, 3, [16, 16], stage=2, block='b',
-                              training=training)
-  x = identity_building_block(x, 3, [16, 16], stage=2, block='c',
-                              training=training)
-  x = identity_building_block(x, 3, [16, 16], stage=2, block='d',
-                              training=training)
-  x = identity_building_block(x, 3, [16, 16], stage=2, block='e',
-                              training=training)
-  x = identity_building_block(x, 3, [16, 16], stage=2, block='f',
-                              training=training)
-  x = identity_building_block(x, 3, [16, 16], stage=2, block='g',
-                              training=training)
-  x = identity_building_block(x, 3, [16, 16], stage=2, block='h',
-                              training=training)
-  x = identity_building_block(x, 3, [16, 16], stage=2, block='i',
-                              training=training)
-
-  x = conv_building_block(x, 3, [32, 32], stage=3, block='a',
-                          training=training)
-  x = identity_building_block(x, 3, [32, 32], stage=3, block='b',
-                              training=training)
-  x = identity_building_block(x, 3, [32, 32], stage=3, block='c',
-                              training=training)
-  x = identity_building_block(x, 3, [32, 32], stage=3, block='d',
-                              training=training)
-  x = identity_building_block(x, 3, [32, 32], stage=3, block='e',
-                              training=training)
-  x = identity_building_block(x, 3, [32, 32], stage=3, block='f',
-                              training=training)
-  x = identity_building_block(x, 3, [32, 32], stage=3, block='g',
-                              training=training)
-  x = identity_building_block(x, 3, [32, 32], stage=3, block='h',
-                              training=training)
-  x = identity_building_block(x, 3, [32, 32], stage=3, block='i',
-                              training=training)
-
-  x = conv_building_block(x, 3, [64, 64], stage=4, block='a',
-                          training=training)
-  x = identity_building_block(x, 3, [64, 64], stage=4, block='b',
-                              training=training)
-  x = identity_building_block(x, 3, [64, 64], stage=4, block='c',
-                              training=training)
-  x = identity_building_block(x, 3, [64, 64], stage=4, block='d',
-                              training=training)
-  x = identity_building_block(x, 3, [64, 64], stage=4, block='e',
-                              training=training)
-  x = identity_building_block(x, 3, [64, 64], stage=4, block='f',
-                              training=training)
-  x = identity_building_block(x, 3, [64, 64], stage=4, block='g',
-                              training=training)
-  x = identity_building_block(x, 3, [64, 64], stage=4, block='h',
-                              training=training)
-  x = identity_building_block(x, 3, [64, 64], stage=4, block='i',
-                              training=training)
+  x = resnet_block(x, size=num_blocks, kernel_size=3, filters=[16, 16],
+                   stage=2, conv_strides=(1, 1), training=training)
+
+  x = resnet_block(x, size=num_blocks, kernel_size=3, filters=[32, 32],
+                   stage=3, conv_strides=(2, 2), training=training)
+
+  x = resnet_block(x, size=num_blocks, kernel_size=3, filters=[64, 64],
+                   stage=4, conv_strides=(2, 2), training=training)
 
   x = tf.keras.layers.GlobalAveragePooling2D(name='avg_pool')(x)
   x = tf.keras.layers.Dense(classes, activation='softmax',
@@ -285,3 +278,9 @@ def resnet56(classes=100, training=None):
   model = tf.keras.models.Model(inputs, x, name='resnet56')
 
   return model
+
+
+resnet20 = functools.partial(resnet, num_blocks=3)
+resnet32 = functools.partial(resnet, num_blocks=5)
+resnet56 = functools.partial(resnet, num_blocks=9)
+resnet10 = functools.partial(resnet, num_blocks=110)