INNER CODE UNIT · Python
EpochCheckpoint
Adibvafa/CodonTransformer · pretrain.py:112
class EpochCheckpoint(pl.Callback):
def __init__(self, checkpoint_dir, save_interval):
super().__init__()
self.checkpoint_dir = checkpoint_dir
self.save_interval = save_interval
def on_train_epoch_end(self, trainer, pl_module):
current_epoch = trainer.current_epoch
if current_epoch % self.save_interval == 0 or current_epoch == 0:
checkpoint_path = os.path.join(
self.checkpoint_dir, f"epoch_{current_epoch}.ckpt"
)
trainer.save_checkpoint(checkpoint_path)
print(f"\nCheckpoint saved at {checkpoint_path}\n")
def main(args):
"""Pretrain the CodonTransformer model."""