remove final_size parameter of resnet (tensorflow#5326)

naurril · Taylor Robie · commit 8ff6115343e3 · 2018-09-18T20:06:50.000-07:00
diff --git a/official/resnet/cifar10_main.py b/official/resnet/cifar10_main.py
@@ -184,7 +184,6 @@ def __init__(self, resnet_size, data_format=None, num_classes=_NUM_CLASSES,
         first_pool_stride=None,
         block_sizes=[num_blocks] * 3,
         block_strides=[1, 2, 2],
-        final_size=64,
         resnet_version=resnet_version,
         data_format=data_format,
         dtype=dtype
diff --git a/official/resnet/imagenet_main.py b/official/resnet/imagenet_main.py
@@ -232,10 +232,8 @@ def __init__(self, resnet_size, data_format=None, num_classes=_NUM_CLASSES,
     # For bigger models, we want to use "bottleneck" layers
     if resnet_size < 50:
       bottleneck = False
-      final_size = 512
     else:
       bottleneck = True
-      final_size = 2048
 
     super(ImagenetModel, self).__init__(
         resnet_size=resnet_size,
@@ -248,7 +246,6 @@ def __init__(self, resnet_size, data_format=None, num_classes=_NUM_CLASSES,
         first_pool_stride=2,
         block_sizes=_get_block_sizes(resnet_size),
         block_strides=[1, 2, 2, 2],
-        final_size=final_size,
         resnet_version=resnet_version,
         data_format=data_format,
         dtype=dtype
diff --git a/official/resnet/resnet_model.py b/official/resnet/resnet_model.py
@@ -354,7 +354,7 @@ def __init__(self, resnet_size, bottleneck, num_classes, num_filters,
                kernel_size,
                conv_stride, first_pool_size, first_pool_stride,
                block_sizes, block_strides,
-               final_size, resnet_version=DEFAULT_VERSION, data_format=None,
+               resnet_version=DEFAULT_VERSION, data_format=None,
                dtype=DEFAULT_DTYPE):
     """Creates a model for classifying an image.
 
@@ -376,7 +376,6 @@ def __init__(self, resnet_size, bottleneck, num_classes, num_filters,
         i-th set.
       block_strides: List of integers representing the desired stride size for
         each of the sets of block layers. Should be same length as block_sizes.
-      final_size: The expected size of the model after the second pooling.
       resnet_version: Integer representing which version of the ResNet network
         to use. See README for details. Valid values: [1, 2]
       data_format: Input format ('channels_last', 'channels_first', or None).
@@ -422,7 +421,6 @@ def __init__(self, resnet_size, bottleneck, num_classes, num_filters,
     self.first_pool_stride = first_pool_stride
     self.block_sizes = block_sizes
     self.block_strides = block_strides
-    self.final_size = final_size
     self.dtype = dtype
     self.pre_activation = resnet_version == 2
 
@@ -542,7 +540,7 @@ def __call__(self, inputs, training):
       inputs = tf.reduce_mean(inputs, axes, keepdims=True)
       inputs = tf.identity(inputs, 'final_reduce_mean')
 
-      inputs = tf.reshape(inputs, [-1, self.final_size])
+      inputs = tf.squeeze(inputs, axes)
       inputs = tf.layers.dense(inputs=inputs, units=self.num_classes)
       inputs = tf.identity(inputs, 'final_dense')
       return inputs
