Save and Check Every N Steps

[rentainhe]

Hi! I've defined a callback like this:

class CheckpointEveryNSteps(pl.Callback):
    """
    Save a checkpoint every N steps, instead of Lightning's default that checkpoints
    based on validation loss.
    """

    def __init__(
        self,
        save_step_frequency=200,
        prefix="",
        use_modelcheckpoint_filename=False,
    ):
        self.save_step_frequency = save_step_frequency
        self.prefix = prefix
        self.use_modelcheckpoint_filename = use_modelcheckpoint_filename

    def on_batch_end(self, trainer: pl.Trainer, _):
        """ Check if we should save a checkpoint after every train batch """
        epoch = trainer.current_epoch
        global_step = trainer.global_step
        if global_step % self.save_step_frequency == 0:
            if self.use_modelcheckpoint_filename:
                filename = trainer.checkpoint_callback.filename
            else:
                filename = f"step={global_step}.ckpt"
            ckpt_path = os.path.join(trainer.checkpoint_callback.dirpath, filename)
            trainer.save_checkpoint(ckpt_path)
            trainer.test(model=trainer.model, dataloaders=trainer.val_dataloaders, ckpt_path=ckpt_path)  # something wrong here

Something went wrong:

  File "/home/rentianhe/anaconda3/envs/oneflow-dev-gcc7/lib/python3.6/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py", line 161, in start_training
    self._results = trainer.run_stage()
  File "/home/rentianhe/anaconda3/envs/oneflow-dev-gcc7/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py", line 995, in run_stage
    return self._run_train()
  File "/home/rentianhe/anaconda3/envs/oneflow-dev-gcc7/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py", line 1044, in _run_train
    self.fit_loop.run()
  File "/home/rentianhe/anaconda3/envs/oneflow-dev-gcc7/lib/python3.6/site-packages/pytorch_lightning/loops/base.py", line 111, in run
    self.advance(*args, **kwargs)
  File "/home/rentianhe/anaconda3/envs/oneflow-dev-gcc7/lib/python3.6/site-packages/pytorch_lightning/loops/fit_loop.py", line 200, in advance
    epoch_output = self.epoch_loop.run(train_dataloader)
  File "/home/rentianhe/anaconda3/envs/oneflow-dev-gcc7/lib/python3.6/site-packages/pytorch_lightning/loops/base.py", line 111, in run
    self.advance(*args, **kwargs)
  File "/home/rentianhe/anaconda3/envs/oneflow-dev-gcc7/lib/python3.6/site-packages/pytorch_lightning/loops/epoch/training_epoch_loop.py", line 153, in advance
    self.trainer.logger_connector.on_batch_end()
  File "/home/rentianhe/anaconda3/envs/oneflow-dev-gcc7/lib/python3.6/site-packages/pytorch_lightning/trainer/connectors/logger_connector/logger_connector.py", line 266, in on_batch_end
    assert not self._epoch_end_reached
AssertionError

I have no idea how to solve this problem

[rohitgr7]

as per the error message, it's coming from the training part.
can you share a reproducible minimal example?
also, you can check ModelCheckpoint callback's every_n_train_steps argument.

also global_step is zero-indexed so this might not be correct

if global_step % self.save_step_frequency == 0: