def tensor_shapes_helper(self, resnet_size, version, with_gpu=False):
"""Checks the tensor shapes after each phase of the ResNet model."""
def reshape(shape):
"""Returns the expected dimensions depending on if a
GPU is being used.
"""
# If a GPU is used for the test, the shape is returned (already in NCHW
# form). When GPU is not used, the shape is converted to NHWC.
if with_gpu:
return shape
return shape[0], shape[2], shape[3], shape[1]
graph = tf.Graph()
with graph.as_default(), self.test_session(use_gpu=with_gpu,
force_gpu=with_gpu):
model = imagenet_main.ImagenetModel(
resnet_size,
data_format='channels_first' if with_gpu else 'channels_last',
version=version)
inputs = tf.random_uniform([1, 224, 224, 3])
output = model(inputs, training=True)
initial_conv = graph.get_tensor_by_name('initial_conv:0')
max_pool = graph.get_tensor_by_name('initial_max_pool:0')
block_layer1 = graph.get_tensor_by_name('block_layer1:0')
block_layer2 = graph.get_tensor_by_name('block_layer2:0')
block_layer3 = graph.get_tensor_by_name('block_layer3:0')
block_layer4 = graph.get_tensor_by_name('block_layer4:0')
avg_pool = graph.get_tensor_by_name('final_avg_pool:0')
dense = graph.get_tensor_by_name('final_dense:0')
self.assertAllEqual(initial_conv.shape, reshape((1, 64, 112, 112)))
self.assertAllEqual(max_pool.shape, reshape((1, 64, 56, 56)))
# The number of channels after each block depends on whether we're
# using the building_block or the bottleneck_block.
if resnet_size < 50:
self.assertAllEqual(block_layer1.shape, reshape(
(1, 64, 56, 56)))
self.assertAllEqual(block_layer2.shape,
reshape((1, 128, 28, 28)))
self.assertAllEqual(block_layer3.shape,
reshape((1, 256, 14, 14)))
self.assertAllEqual(block_layer4.shape, reshape(
(1, 512, 7, 7)))
self.assertAllEqual(avg_pool.shape, reshape((1, 512, 1, 1)))
else:
self.assertAllEqual(block_layer1.shape,
reshape((1, 256, 56, 56)))
self.assertAllEqual(block_layer2.shape,
reshape((1, 512, 28, 28)))
self.assertAllEqual(block_layer3.shape,
reshape((1, 1024, 14, 14)))
self.assertAllEqual(block_layer4.shape, reshape(
(1, 2048, 7, 7)))
self.assertAllEqual(avg_pool.shape, reshape((1, 2048, 1, 1)))
self.assertAllEqual(dense.shape, (1, _LABEL_CLASSES))
self.assertAllEqual(output.shape, (1, _LABEL_CLASSES))
def test_imagenetmodel_shape(self):
batch_size = 135
num_classes = 246
model = imagenet_main.ImagenetModel(50,
data_format='channels_last',
num_classes=num_classes)
fake_input = tf.random_uniform([batch_size, 224, 224, 3])
output = model(fake_input, training=True)
self.assertAllEqual(output.shape, (batch_size, num_classes))
def create_network_resnet(inputs, embedding_dim, is_training, reuse=False):
resnet_size = 50
model = imagenet_main.ImagenetModel(resnet_size, num_classes=embedding_dim)
with tf.variable_scope("", reuse=reuse):
y = model(inputs, training=is_training)
return y
