def test_training_hooks_are_used(self):
chief_hook = test.mock.MagicMock(
wraps=training.SessionRunHook(), spec=training.SessionRunHook)
hook = test.mock.MagicMock(
wraps=training.SessionRunHook(), spec=training.SessionRunHook)
def _model_fn_hooks(features, labels, mode):
_, _ = features, labels
return model_fn_lib.EstimatorSpec(
mode=mode,
loss=constant_op.constant(0.),
train_op=constant_op.constant(0.),
training_chief_hooks=[chief_hook],
training_hooks=[hook])
est = estimator.Estimator(model_fn=_model_fn_hooks)
self.assertFalse(chief_hook.begin.called)
self.assertFalse(hook.begin.called)
est.fit(dummy_input_fn, steps=1)
self.assertTrue(chief_hook.begin.called)
self.assertTrue(hook.begin.called)
def test_chief_only_hook_should_not_be_called_on_non_chief(self):
chief_hook = test.mock.MagicMock(
wraps=training.SessionRunHook(), spec=training.SessionRunHook)
hook = test.mock.MagicMock(
wraps=training.SessionRunHook(), spec=training.SessionRunHook)
def _model_fn_hooks(features, labels, mode):
_, _ = features, labels
return model_fn_lib.EstimatorSpec(
mode=mode,
loss=constant_op.constant(0.),
train_op=constant_op.constant(0.),
training_chief_hooks=[chief_hook],
training_hooks=[hook])
class NonChiefRunConfig(run_config.RunConfig): # pylint: disable=g-wrong-blank-lines
@property
def is_chief(self): # pylint: disable=g-wrong-blank-lines
return False
# Mocking the SessionManager.wait_for_session, so that worker doesn't wait
# for chief.
def get_initialized_session(*args, **kwargs):
scaffold = training.Scaffold().finalize()
sess = session.Session(*args, **kwargs)
sess.run(scaffold.init_op)
return sess
with test.mock.patch.object(
training.SessionManager,
'wait_for_session',
side_effect=get_initialized_session):
est = estimator.Estimator(
model_fn=_model_fn_hooks, config=NonChiefRunConfig())
self.assertFalse(chief_hook.begin.called)
self.assertFalse(hook.begin.called)
est.fit(dummy_input_fn, steps=1)
self.assertFalse(chief_hook.begin.called)
self.assertTrue(hook.begin.called)