| import os | |
| import pytorch_lightning as pl | |
| from pytorch_lightning.callbacks.base import Callback | |
| class IntervalModelCheckpoint(Callback): | |
| """ | |
| Save a checkpoint every N steps, instead of Lightning's default that checkpoints | |
| based on validation loss. | |
| """ | |
| def __init__( | |
| self, | |
| dirpath, | |
| save_intervals, | |
| ): | |
| """ | |
| Args: | |
| save_step_frequency: how often to save in steps | |
| prefix: add a prefix to the name, only used if | |
| use_modelcheckpoint_filename=False | |
| use_modelcheckpoint_filename: just use the ModelCheckpoint callback's | |
| default filename, don't use ours. | |
| """ | |
| self.dirpath = dirpath | |
| self.save_intervals = save_intervals | |
| self.best_val_loss = 1e10 | |
| def on_batch_end(self, trainer: pl.Trainer, _): | |
| """ Check if we should save a checkpoint after every train batch """ | |
| global_step = trainer.global_step | |
| if (global_step + 1) in self.save_intervals: | |
| trainer.run_evaluation() | |
| val_loss = trainer.callback_metrics['val_loss'] | |
| filename = f"steps={global_step+1:05d}-val_loss={val_loss:0.8f}.ckpt" | |
| ckpt_path = os.path.join(self.dirpath, filename) | |
| trainer.save_checkpoint(ckpt_path) | |
| if val_loss < self.best_val_loss: | |
| best_ckpt_path = os.path.join(self.dirpath, 'best.ckpt') | |
| trainer.save_checkpoint(best_ckpt_path) | |
| self.best_val_loss = val_loss | |