| import os |
| import lightning.pytorch as pl |
| from lightning.pytorch.utilities import rank_zero_only |
|
|
|
|
| class CheckpointEveryNSteps(pl.Callback): |
| def __init__( |
| self, |
| checkpoints_dir, |
| save_step_frequency, |
| ) -> None: |
| r"""Save a checkpoint every N steps. |
| |
| Args: |
| checkpoints_dir (str): directory to save checkpoints |
| save_step_frequency (int): save checkpoint every N step |
| """ |
|
|
| self.checkpoints_dir = checkpoints_dir |
| self.save_step_frequency = save_step_frequency |
|
|
| @rank_zero_only |
| def on_train_batch_end(self, *args, **kwargs) -> None: |
| r"""Save a checkpoint every N steps.""" |
|
|
| trainer = args[0] |
| global_step = trainer.global_step |
|
|
| if global_step == 1 or global_step % self.save_step_frequency == 0: |
|
|
| ckpt_path = os.path.join( |
| self.checkpoints_dir, |
| "step={}.ckpt".format(global_step)) |
| trainer.save_checkpoint(ckpt_path) |
| print("Save checkpoint to {}".format(ckpt_path)) |
|
|
