Spaces:
Running
on
A10G
Running
on
A10G
| 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 | |
| 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)) | |