From 125a5b41ff30e3960cf632326d5022bac1c68fc3 Mon Sep 17 00:00:00 2001 From: Stephen Roller Date: Sat, 29 Aug 2020 18:17:59 -0400 Subject: [PATCH] Fix more tests. --- parlai/nn/lr_scheduler.py | 4 ++-- tests/test_lr_schedulers.py | 2 +- tests/test_transformers.py | 1 + 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/parlai/nn/lr_scheduler.py b/parlai/nn/lr_scheduler.py index 7d72b39f8a0..a1122c99d16 100644 --- a/parlai/nn/lr_scheduler.py +++ b/parlai/nn/lr_scheduler.py @@ -60,7 +60,7 @@ def _init_warmup_scheduler(self, optimizer, states): self.warmup_scheduler = optim.lr_scheduler.LambdaLR( optimizer, self._warmup_lr ) - if 'warmup_scheduler' in states: + if states.get('warmup_scheduler'): self.warmup_scheduler.load_state_dict(states['warmup_scheduler']) else: self.warmup_scheduler = None @@ -307,7 +307,7 @@ def step(self, num_steps): Override this method to override the behavior for training schedulers. """ - self._number_training_updates += 1 + self._number_training_updates = num_steps if self._is_lr_warming_up(): self.warmup_scheduler.step() else: diff --git a/tests/test_lr_schedulers.py b/tests/test_lr_schedulers.py index 439f198f763..22b0afacee1 100644 --- a/tests/test_lr_schedulers.py +++ b/tests/test_lr_schedulers.py @@ -96,7 +96,7 @@ def test_invsqrt(self): steps = self._run_pass( lr_scheduler='invsqrt', warmup_updates=50, invsqrt_lr_decay_gamma=1 ) - self.assertAlmostEquals(steps[-1], 0.0324272) + self.assertAlmostEquals(steps[-1], 0.0324443) # decay very slowly steps = self._run_pass( diff --git a/tests/test_transformers.py b/tests/test_transformers.py index 44fb2338aa6..4b548345c2e 100644 --- a/tests/test_transformers.py +++ b/tests/test_transformers.py @@ -797,6 +797,7 @@ def _test_learning_rate_resuming(self, args): init_model=os.path.join(tmpdir, 'model'), model_file=os.path.join(tmpdir, 'newmodel2'), lr_scheduler='reduceonplateau', + log_every_n_secs=0.001, **args, ) )