def test_make_best_model_export_strategy(self):
export_dir_base = tempfile.mkdtemp() + "export/"
gfile.MkDir(export_dir_base)
test_estimator = TestEstimator()
export_strategy = saved_model_export_utils.make_best_model_export_strategy(
serving_input_fn=None, exports_to_keep=3, compare_fn=None)
self.assertNotEqual("",
export_strategy.export(test_estimator, export_dir_base,
"fake_ckpt_0", {"loss": 100}))
self.assertNotEqual("", test_estimator.last_exported_dir)
self.assertNotEqual("", test_estimator.last_exported_checkpoint)
self.assertEqual("",
export_strategy.export(test_estimator, export_dir_base,
"fake_ckpt_1", {"loss": 101}))
self.assertEqual(test_estimator.last_exported_dir,
os.path.join(export_dir_base, "fake_ckpt_0"))
self.assertNotEqual("",
export_strategy.export(test_estimator, export_dir_base,
"fake_ckpt_2", {"loss": 10}))
self.assertEqual(test_estimator.last_exported_dir,
os.path.join(export_dir_base, "fake_ckpt_2"))
self.assertEqual("",
export_strategy.export(test_estimator, export_dir_base,
"fake_ckpt_3", {"loss": 20}))
self.assertEqual(test_estimator.last_exported_dir,
os.path.join(export_dir_base, "fake_ckpt_2"))
def test_make_best_model_export_strategy_exceptions(self):
export_dir_base = tempfile.mkdtemp() + "export/"
test_estimator = TestEstimator()
export_strategy = saved_model_export_utils.make_best_model_export_strategy(
serving_input_fn=None, exports_to_keep=3, compare_fn=None)
with self.assertRaises(ValueError):
export_strategy.export(test_estimator, export_dir_base, "", {"loss": 200})
with self.assertRaises(ValueError):
export_strategy.export(test_estimator, export_dir_base, "fake_ckpt_1",
None)
def test_make_best_model_export_strategy_with_preemption(self):
model_dir = self.get_temp_dir()
eval_dir_base = os.path.join(model_dir, "eval_continuous")
core_estimator._write_dict_to_summary(eval_dir_base, {"loss": 50}, 1)
core_estimator._write_dict_to_summary(eval_dir_base, {"loss": 60}, 2)
test_estimator = TestEstimator()
export_strategy = saved_model_export_utils.make_best_model_export_strategy(
serving_input_fn=None,
exports_to_keep=3,
model_dir=model_dir,
event_file_pattern="eval_continuous/*.tfevents.*",
compare_fn=None)
export_dir_base = os.path.join(self.get_temp_dir(), "export")
self.assertEqual("",
export_strategy.export(test_estimator, export_dir_base,
"fake_ckpt_0", {
"loss": 100
}))
self.assertEqual("", test_estimator.last_exported_dir)
self.assertEqual("", test_estimator.last_exported_checkpoint)
self.assertNotEqual("",
export_strategy.export(test_estimator, export_dir_base,
"fake_ckpt_2", {
"loss": 10
}))
self.assertEqual(test_estimator.last_exported_dir,
os.path.join(export_dir_base, "fake_ckpt_2"))
self.assertEqual("",
export_strategy.export(test_estimator, export_dir_base,
"fake_ckpt_3", {
"loss": 20
}))
self.assertEqual(test_estimator.last_exported_dir,
os.path.join(export_dir_base, "fake_ckpt_2"))