|
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)) |
|
|
