def test_torch_model_on_epoch_end_lr_wait(self):
model = fe.build(model_fn=MultiLayerTorchModel, optimizer_fn='adam')
lr_on_plateau = ReduceLROnPlateau(model=model, metric='loss')
lr_on_plateau.system = sample_system_object()
lr_on_plateau.best = 12
lr_on_plateau.on_epoch_end(data=self.data)
with self.subTest('Check value of wait'):
self.assertEqual(lr_on_plateau.wait, 0)
with self.subTest('Check value of best'):
self.assertEqual(lr_on_plateau.best, 10)
def test_torch_model_on_epoch_end_reduce_lr(self):
model = fe.build(model_fn=MultiLayerTorchModel, optimizer_fn='adam')
model_name = model.model_name + '_lr'
lr_on_plateau = ReduceLROnPlateau(model=model, metric='loss')
lr_on_plateau.system = sample_system_object()
lr_on_plateau.best = 5
lr_on_plateau.wait = 11
with patch('sys.stdout', new=StringIO()) as fake_stdout:
lr_on_plateau.on_epoch_end(data=self.data)
log = fake_stdout.getvalue().strip()
self.assertEqual(log, self.torch_expected_msg)
with self.subTest('Check learning rate in data'):
self.assertTrue(math.isclose(self.data[model_name], 0.000100000005, rel_tol=1e-3))