diff --git a/pytorch_lightning/trainer/__init__.py b/pytorch_lightning/trainer/__init__.py index 8cb144cb..554b34f8 100644 --- a/pytorch_lightning/trainer/__init__.py +++ b/pytorch_lightning/trainer/__init__.py @@ -30,6 +30,7 @@ This is the basic use of the trainer: trainer = Trainer() trainer.fit(model) + -------- Best Practices @@ -59,6 +60,14 @@ So you can run it like so:distributed_backend $ python main.py --gpus 2 + +.. note:: + If you want to stop a training run early, you can press "Ctrl + C" on your keyboard. + The trainer will catch the `KeyboardInterrupt` and attempt a graceful shutdown, including + running callbacks such as `on_train_end`. The trainer object will also set an attribute + `interrupted` to `True` in such cases. If you have a callback which shuts down compute + resources, for example, you can conditionally run the shutdown logic for only uninterrupted runs. + ------------ Testing diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py index dca3d539..7b0568e7 100644 --- a/pytorch_lightning/trainer/trainer.py +++ b/pytorch_lightning/trainer/trainer.py @@ -364,6 +364,7 @@ class Trainer( self.global_step = 0 self.current_epoch = 0 self.total_batches = 0 + self.interrupted = False # configure logger self.configure_logger(logger) diff --git a/pytorch_lightning/trainer/training_loop.py b/pytorch_lightning/trainer/training_loop.py index 071a2f94..ba974dac 100644 --- a/pytorch_lightning/trainer/training_loop.py +++ b/pytorch_lightning/trainer/training_loop.py @@ -203,6 +203,7 @@ class TrainerTrainLoopMixin(ABC): use_amp: bool track_grad_norm: ... model: LightningModule + interrupted: bool running_loss: ... training_tqdm_dict: ... reduce_lr_on_plateau_scheduler: ... @@ -387,6 +388,7 @@ class TrainerTrainLoopMixin(ABC): except KeyboardInterrupt: log.info('Detected KeyboardInterrupt, attempting graceful shutdown...') + self.interrupted = True self.run_training_teardown() def run_training_epoch(self): diff --git a/tests/trainer/test_trainer.py b/tests/trainer/test_trainer.py index ddf79cb2..f80870ec 100644 --- a/tests/trainer/test_trainer.py +++ b/tests/trainer/test_trainer.py @@ -12,6 +12,7 @@ from pytorch_lightning.callbacks import ( EarlyStopping, ModelCheckpoint, ) +from pytorch_lightning import Callback from pytorch_lightning.core.lightning import load_hparams_from_tags_csv from pytorch_lightning.trainer.logging import TrainerLoggingMixin from pytorch_lightning.utilities.exceptions import MisconfigurationException @@ -630,3 +631,33 @@ def test_nan_params_detection(tmpdir): # after aborting the training loop, model still has nan-valued params params = torch.cat([param.view(-1) for param in model.parameters()]) assert not torch.isfinite(params).all() + + +def test_trainer_interrupted_flag(tmpdir): + """Test the flag denoting that a user interrupted training.""" + + model = DictHparamsModel({'in_features': 28 * 28, 'out_features': 10}) + + class InterruptCallback(Callback): + def __init__(self): + super().__init__() + + def on_batch_start(self, trainer, pl_module): + raise KeyboardInterrupt + + interrupt_callback = InterruptCallback() + + trainer_options = { + 'callbacks': [interrupt_callback], + 'max_epochs': 1, + 'val_percent_check': 0.1, + 'train_percent_check': 0.2, + 'progress_bar_refresh_rate': 0, + 'logger': False, + 'default_save_path': tmpdir, + } + + trainer = Trainer(**trainer_options) + assert not trainer.interrupted + trainer.fit(model) + assert trainer.interrupted