wehub-resource-sync
207 lines
8.4 KiB
Python
207 lines
8.4 KiB
Python
import logging
|
|
import os
|
|
import posixpath
|
|
|
|
import mlflow
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.utils.autologging_utils import (
|
|
ExceptionSafeAbstractClass,
|
|
)
|
|
from mlflow.utils.file_utils import TempDir
|
|
from mlflow.utils.mlflow_tags import LATEST_CHECKPOINT_ARTIFACT_TAG_KEY
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
|
|
_CHECKPOINT_DIR = "checkpoints"
|
|
_CHECKPOINT_METRIC_FILENAME = "checkpoint_metrics.json"
|
|
_CHECKPOINT_MODEL_FILENAME = "checkpoint"
|
|
_LATEST_CHECKPOINT_PREFIX = "latest_"
|
|
_CHECKPOINT_EPOCH_PREFIX = "epoch_"
|
|
_CHECKPOINT_GLOBAL_STEP_PREFIX = "global_step_"
|
|
_WEIGHT_ONLY_CHECKPOINT_SUFFIX = ".weights"
|
|
|
|
|
|
class MlflowModelCheckpointCallbackBase(metaclass=ExceptionSafeAbstractClass):
|
|
"""Callback base class for automatic model checkpointing to MLflow.
|
|
|
|
You must implement "save_checkpoint" method to save the model as the checkpoint file.
|
|
and you must call `check_and_save_checkpoint_if_needed` method in relevant
|
|
callback events to trigger automatic checkpointing.
|
|
|
|
Args:
|
|
checkpoint_file_suffix: checkpoint file suffix.
|
|
monitor: In automatic model checkpointing, the metric name to monitor if
|
|
you set `model_checkpoint_save_best_only` to True.
|
|
save_best_only: If True, automatic model checkpointing only saves when
|
|
the model is considered the "best" model according to the quantity
|
|
monitored and previous checkpoint model is overwritten.
|
|
mode: one of {"min", "max"}. In automatic model checkpointing,
|
|
if save_best_only=True, the decision to overwrite the current save file is made
|
|
based on either the maximization or the minimization of the monitored quantity.
|
|
save_weights_only: In automatic model checkpointing, if True, then
|
|
only the model's weights will be saved. Otherwise, the optimizer states,
|
|
lr-scheduler states, etc are added in the checkpoint too.
|
|
save_freq: `"epoch"` or integer. When using `"epoch"`, the callback
|
|
saves the model after each epoch. When using integer, the callback
|
|
saves the model at end of this many batches. Note that if the saving isn't
|
|
aligned to epochs, the monitored metric may potentially be less reliable (it
|
|
could reflect as little as 1 batch, since the metrics get reset
|
|
every epoch). Defaults to `"epoch"`.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
checkpoint_file_suffix,
|
|
monitor,
|
|
mode,
|
|
save_best_only,
|
|
save_weights_only,
|
|
save_freq,
|
|
):
|
|
self.checkpoint_file_suffix = checkpoint_file_suffix
|
|
self.monitor = monitor
|
|
self.mode = mode
|
|
self.save_best_only = save_best_only
|
|
self.save_weights_only = save_weights_only
|
|
self.save_freq = save_freq
|
|
self.last_monitor_value = None
|
|
|
|
self.mlflow_tracking_uri = mlflow.get_tracking_uri()
|
|
|
|
if self.save_best_only:
|
|
if self.monitor is None:
|
|
raise MlflowException(
|
|
"If checkpoint 'save_best_only' config is set to True, you need to set "
|
|
"'monitor' config as well."
|
|
)
|
|
if self.mode not in ["min", "max"]:
|
|
raise MlflowException(
|
|
"If checkpoint 'save_best_only' config is set to True, you need to set "
|
|
"'mode' config and available modes includes 'min' and 'max', but you set "
|
|
f"'mode' to '{self.mode}'."
|
|
)
|
|
|
|
def _is_new_checkpoint_better(self, new_monitor_value):
|
|
if self.last_monitor_value is None:
|
|
return True
|
|
|
|
if self.mode == "min":
|
|
return new_monitor_value < self.last_monitor_value
|
|
|
|
return new_monitor_value > self.last_monitor_value
|
|
|
|
def save_checkpoint(self, filepath: str):
|
|
raise NotImplementedError()
|
|
|
|
def check_and_save_checkpoint_if_needed(self, current_epoch, global_step, metric_dict):
|
|
# For distributed model training, trainer workers need to use the driver process
|
|
# mlflow_tracking_uri.
|
|
# Note that `self.mlflow_tracking_uri` value is assigned in the driver process
|
|
# then it is pickled to trainer workers.
|
|
mlflow.set_tracking_uri(self.mlflow_tracking_uri)
|
|
|
|
if self.save_best_only:
|
|
if self.monitor not in metric_dict:
|
|
_logger.warning(
|
|
"Checkpoint logging is skipped, because checkpoint 'save_best_only' config is "
|
|
"True, it requires to compare the monitored metric value, but the provided "
|
|
"monitored metric value is not available."
|
|
)
|
|
return
|
|
|
|
new_monitor_value = metric_dict[self.monitor]
|
|
if not self._is_new_checkpoint_better(new_monitor_value):
|
|
# Current checkpoint is worse than last saved checkpoint,
|
|
# so skip checkpointing.
|
|
return
|
|
|
|
self.last_monitor_value = new_monitor_value
|
|
|
|
suffix = self.checkpoint_file_suffix
|
|
|
|
if self.save_best_only:
|
|
if self.save_weights_only:
|
|
checkpoint_model_filename = (
|
|
f"{_LATEST_CHECKPOINT_PREFIX}{_CHECKPOINT_MODEL_FILENAME}"
|
|
f"{_WEIGHT_ONLY_CHECKPOINT_SUFFIX}{suffix}"
|
|
)
|
|
else:
|
|
checkpoint_model_filename = (
|
|
f"{_LATEST_CHECKPOINT_PREFIX}{_CHECKPOINT_MODEL_FILENAME}{suffix}"
|
|
)
|
|
checkpoint_metrics_filename = (
|
|
f"{_LATEST_CHECKPOINT_PREFIX}{_CHECKPOINT_METRIC_FILENAME}"
|
|
)
|
|
checkpoint_artifact_dir = _CHECKPOINT_DIR
|
|
else:
|
|
if self.save_freq == "epoch":
|
|
sub_dir_name = f"{_CHECKPOINT_EPOCH_PREFIX}{current_epoch}"
|
|
else:
|
|
sub_dir_name = f"{_CHECKPOINT_GLOBAL_STEP_PREFIX}{global_step}"
|
|
|
|
if self.save_weights_only:
|
|
checkpoint_model_filename = (
|
|
f"{_CHECKPOINT_MODEL_FILENAME}{_WEIGHT_ONLY_CHECKPOINT_SUFFIX}{suffix}"
|
|
)
|
|
else:
|
|
checkpoint_model_filename = f"{_CHECKPOINT_MODEL_FILENAME}{suffix}"
|
|
checkpoint_metrics_filename = _CHECKPOINT_METRIC_FILENAME
|
|
checkpoint_artifact_dir = f"{_CHECKPOINT_DIR}/{sub_dir_name}"
|
|
|
|
mlflow.set_tag(
|
|
LATEST_CHECKPOINT_ARTIFACT_TAG_KEY,
|
|
f"{checkpoint_artifact_dir}/{checkpoint_model_filename}",
|
|
)
|
|
|
|
mlflow.log_dict(
|
|
{**metric_dict, "epoch": current_epoch, "global_step": global_step},
|
|
f"{checkpoint_artifact_dir}/{checkpoint_metrics_filename}",
|
|
)
|
|
|
|
with TempDir() as tmp_dir:
|
|
tmp_model_save_path = os.path.join(tmp_dir.path(), checkpoint_model_filename)
|
|
self.save_checkpoint(tmp_model_save_path)
|
|
mlflow.log_artifact(tmp_model_save_path, checkpoint_artifact_dir)
|
|
|
|
|
|
def download_checkpoint_artifact(run_id=None, epoch=None, global_step=None, dst_path=None):
|
|
from mlflow.client import MlflowClient
|
|
from mlflow.utils.mlflow_tags import LATEST_CHECKPOINT_ARTIFACT_TAG_KEY
|
|
|
|
client = MlflowClient()
|
|
|
|
if run_id is None:
|
|
run = mlflow.active_run()
|
|
if run is None:
|
|
raise MlflowException(
|
|
"There is no active run, please provide the 'run_id' argument for "
|
|
"'load_checkpoint' invocation."
|
|
)
|
|
run_id = run.info.run_id
|
|
else:
|
|
run = client.get_run(run_id)
|
|
|
|
latest_checkpoint_artifact_path = run.data.tags.get(LATEST_CHECKPOINT_ARTIFACT_TAG_KEY)
|
|
if latest_checkpoint_artifact_path is None:
|
|
raise MlflowException("There is no logged checkpoint artifact in the current run.")
|
|
|
|
checkpoint_filename = posixpath.basename(latest_checkpoint_artifact_path)
|
|
|
|
if epoch is not None and global_step is not None:
|
|
raise MlflowException(
|
|
"Only one of 'epoch' and 'global_step' can be set for 'load_checkpoint'."
|
|
)
|
|
elif global_step is not None:
|
|
checkpoint_artifact_path = (
|
|
f"{_CHECKPOINT_DIR}/{_CHECKPOINT_GLOBAL_STEP_PREFIX}{global_step}/{checkpoint_filename}"
|
|
)
|
|
elif epoch is not None:
|
|
checkpoint_artifact_path = (
|
|
f"{_CHECKPOINT_DIR}/{_CHECKPOINT_EPOCH_PREFIX}{epoch}/{checkpoint_filename}"
|
|
)
|
|
else:
|
|
checkpoint_artifact_path = latest_checkpoint_artifact_path
|
|
|
|
return client.download_artifacts(run_id, checkpoint_artifact_path, dst_path=dst_path)
|