def _test_cifar10model_shape(self, resnet_version):
batch_size = 135
num_classes = 246
model = cifar10_main.Cifar10Model(32, data_format='channels_last',
num_classes=num_classes,
resnet_version=resnet_version)
fake_input = tf.random_uniform([batch_size, _HEIGHT, _WIDTH, _NUM_CHANNELS])
output = model(fake_input, training=True)
self.assertAllEqual(output.shape, (batch_size, num_classes))
def test_cifar10model_shape(self):
batch_size = 135
num_classes = 246
for version in (1, 2):
model = cifar10_main.Cifar10Model(
32,
data_format="channels_last",
num_classes=num_classes,
version=version,
)
fake_input = tf.random_uniform(
[batch_size, _HEIGHT, _WIDTH, _NUM_CHANNELS])
output = model(fake_input, training=True)
self.assertAllEqual(output.shape, (batch_size, num_classes))