Method save

pythia/utils/checkpoint.py:204–233

(self, iteration, update_best=False)

        }

    def save(self, iteration, update_best=False):
        # Only save in main process
        if not is_main_process():
            return

        ckpt_filepath = os.path.join(
            self.models_foldername, "model_%d.ckpt" % iteration
        )
        best_ckpt_filepath = os.path.join(
            self.ckpt_foldername, self.ckpt_prefix + "best.ckpt"
        )

        best_iteration = self.trainer.early_stopping.best_monitored_iteration
        best_metric = self.trainer.early_stopping.best_monitored_value

        ckpt = {
            "model": self.trainer.model.state_dict(),
            "optimizer": self.trainer.optimizer.state_dict(),
            "best_iteration": best_iteration,
            "best_metric_value": best_metric,
            "config": self.config,
        }

        git_metadata_dict = self._get_vcs_fields()
        ckpt.update(git_metadata_dict)

        torch.save(ckpt, ckpt_filepath)

        if update_best:
            torch.save(ckpt, best_ckpt_filepath)

    def restore(self):
        self.trainer.writer.write("Restoring checkpoint")