def test_model_registration(self):
@registry.register_model
class MyModel1(Model):
pass
model = registry.model("my_model1")
self.assertTrue(model is MyModel1)
def test_named_registration(self):
@registry.register_model("model2")
class MyModel1(Model):
pass
model = registry.model("model2")
self.assertTrue(model is MyModel1)
'train_batch_size': 32,
'eval_batch_size': 1,
'data_format': 'channels_last',
'optimizer': 'sgd',
'learning_rate': 1e-3,
'momentum': 0.96,
'weight_decay': 2e-4,
'gradient_clipping': None,
'lr_decay_steps': None,
'init_dir': None, # for pre-trained models
'sync': False
}
# Model hyperparameters and definition
model_hparams = {}
model = registry.model('single_char_baseline')(num_classes=task.num_classes,
**model_hparams)
# Unique job_id
experiment_id = 'example'
shape = re.sub(r'([,])', '_', re.sub(r'([() ])', '', str(args['shape_in'])))
job_id = os.path.join(experiment_id, shape, model.name)
task.job_id = job_id # Used to check for first model initialization.
job_dir = os.path.join(task.task_log_dir, job_id)
# TensorFlow Configuration
sess_config = tf.ConfigProto(
allow_soft_placement=True,
log_device_placement=False,
intra_op_parallelism_threads=0,
gpu_options=tf.GPUOptions(force_gpu_compatible=True))
config = tf.estimator.RunConfig(session_config=sess_config,
def test_request_unprovided_model(self):
with self.assertRaisesRegex(LookupError, "never registered"):
_ = registry.model("not_provided")
def test_access_provided_model(self):
model = registry.model("single_char_baseline")
self.assertTrue(model is SingleCharBaseline)