From efd750565e62734d9eaa6a3817243dcb970950e1 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Sun, 21 Apr 2019 12:29:48 -0400 Subject: [PATCH] added early epoch stopping hook --- src/pytorch-lightning/models/trainer.py | 3 ++- src/pytorch-lightning/root_module/hooks.py | 2 -- src/pytorch-lightning/root_module/root_module.py | 1 + 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/pytorch-lightning/models/trainer.py b/src/pytorch-lightning/models/trainer.py index 1088225f..7d6f43a8 100644 --- a/src/pytorch-lightning/models/trainer.py +++ b/src/pytorch-lightning/models/trainer.py @@ -267,7 +267,8 @@ class Trainer(TrainerIO): break # give model a chance to end epoch early - if self.model.should_stop_epoch(): + if self.model.should_stop_epoch: + self.model.should_stop_epoch = False break # --------------- diff --git a/src/pytorch-lightning/root_module/hooks.py b/src/pytorch-lightning/root_module/hooks.py index e5a62f01..0d903f66 100644 --- a/src/pytorch-lightning/root_module/hooks.py +++ b/src/pytorch-lightning/root_module/hooks.py @@ -19,5 +19,3 @@ class ModelHooks(torch.nn.Module): def on_post_performance_check(self): pass - def should_stop_epoch(self): - return False diff --git a/src/pytorch-lightning/root_module/root_module.py b/src/pytorch-lightning/root_module/root_module.py index ab49a9a1..b12a2087 100644 --- a/src/pytorch-lightning/root_module/root_module.py +++ b/src/pytorch-lightning/root_module/root_module.py @@ -24,6 +24,7 @@ class RootModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks): self.overfit = hparams.overfit self.gradient_clip = hparams.gradient_clip self.num = 2 + self.should_stop_epoch = False # track if gpu was requested for checkpointing self.on_gpu = False