Pytorch-lightning: support graceful training cleanup after Keyboard Interrupt

[jeremyjordan]

🚀 Feature

In the training loop, we should catch KeyboardInterrupt exceptions and enter a graceful termination (ie. run on_training_end, print profiler report, etc).

Motivation

Sometimes I'll train a model and just let it run until something interesting happens or my loss has converged. Exiting via "Ctrl + C" can be easier than waiting for "EarlyStopping" to trigger in some cases. However, I may still want certain things to occur (eg. printing out the profiler summary) after stopping training.

Salesforce's AWD-LSTM has an example of this feature implemented here.

Pitch

Figure out what actions we want to run after a KeyboardInterrupt. Then adjust the training loop to catch this exception and ensure the aforementioned actions still run afterwards.

[williamFalcon]

awesome, let's do it

[williamFalcon]

@jeremyjordan can we get it into this release? :)

[jeremyjordan]

when’s the release? planning to get to it this weekend

[williamFalcon]

feb 21

[jeremyjordan]

sure, shouldn't be a problem then :)

[lezwon]

Hey guys, I am currently training a model on Sagemaker and have added a notebook shutdown code within the on_training_end method. There were times where I had to manually cancel my model training because something was wrong. However, If I do that, the notebook shuts down. Is there any way to make this feature optional? I don't want my notebook shutting down after a keyboard interrupt, only after successful training completion.

[jeremyjordan]

@lezwon Hmm yeah, that's not ideal. Do you mind creating a new issue (you can reference this one) with the problem and I can draft up a PR to address it?