def test_create_model(self):
"""Tests creating model."""
model = train.create_model(half_precision=False)
params, state = train.initialized(random.PRNGKey(0), 224, model)
variables = {'params': params, **state}
x = random.normal(random.PRNGKey(1), (8, 224, 224, 3))
y = model.apply(variables, x, train=False)
self.assertEqual(y.shape, (8, 1000))
def test_create_model(self):
"""Tests creating model."""
model = train.create_model(model_cls=models._ResNet1, half_precision=False) # pylint: disable=protected-access
params, state = train.initialized(random.PRNGKey(0), 224, model)
variables = {'params': params, **state}
x = random.normal(random.PRNGKey(1), (8, 224, 224, 3))
y = model.apply(variables, x, train=False)
self.assertEqual(y.shape, (8, 1000))
def test_create_model(self):
variables = train.initialized(random.PRNGKey(0), 224)
x = random.normal(random.PRNGKey(1), (8, 224, 224, 3))
y = train.model(train=False).apply(variables, x)
self.assertEqual(y.shape, (8, 1000))