def _getCheckpointCallback(self, version: int = None, stage=None):
model_dir = self.paths.get("model_dir", f'model/{DATE}')
model_dir = os.path.join(model_dir, self.name)
if version or version == 0:
version = f"version_{version}" if isinstance(version,
int) else version
model_dir = os.path.join(model_dir, version)
exist = os.path.exists(model_dir) and os.listdir(model_dir)
if stage == 'test' and not exist:
raise FileNotFoundError(model_dir)
elif stage is None and exist:
raise FileExistsError(model_dir)
self.model_dir = model_dir
checkpoint_callback = ModelCheckpoint(
monitor="err/B-M/validation",
dirpath=model_dir,
filename="best",
save_last=True,
mode="min",
)
checkpoint_callback.FILE_EXTENSION = ".pth"
checkpoint_callback.best_model_path = model_dir
checkpoint_callback.CHECKPOINT_NAME_LAST = "latest"
return checkpoint_callback
