def test_serialize_config(self): """Tests functionality for serializing data.""" config = base_configs.ExperimentConfig() model_dir = self.create_tempdir().full_path classifier_trainer.serialize_config(params=config, model_dir=model_dir) saved_params_path = os.path.join(model_dir, 'params.yaml') self.assertTrue(os.path.exists(saved_params_path)) tf.io.gfile.rmtree(model_dir)
def test_get_model_size(self, model, model_name, expected):
config = base_configs.ExperimentConfig(
model_name=model,
model=base_configs.ModelConfig(model_params={
'model_name': model_name,
}, ))
size = classifier_trainer.get_image_size_from_model(config)
self.assertEqual(size, expected)
def test_initialize(self, dtype):
config = base_configs.ExperimentConfig(
runtime=base_configs.RuntimeConfig(
run_eagerly=False,
enable_xla=False,
per_gpu_thread_count=1,
gpu_thread_mode='gpu_private',
num_gpus=1,
dataset_num_private_threads=1,
),
train_dataset=dataset_factory.DatasetConfig(dtype=dtype),
model=base_configs.ModelConfig(),
)
class EmptyClass:
pass
fake_ds_builder = EmptyClass()
fake_ds_builder.dtype = dtype
fake_ds_builder.config = EmptyClass()
classifier_trainer.initialize(config, fake_ds_builder)
def test_get_loss_scale(self, loss_scale, dtype, expected):
config = base_configs.ExperimentConfig(
runtime=base_configs.RuntimeConfig(loss_scale=loss_scale),
train_dataset=dataset_factory.DatasetConfig(dtype=dtype))
ls = classifier_trainer.get_loss_scale(config, fp16_default=128)
self.assertEqual(ls, expected)