| |
| from typing import Dict, Any, Optional, List |
| import os |
| import attr |
| from mlagents.trainers.training_status import GlobalTrainingStatus, StatusType |
| from mlagents_envs.logging_util import get_logger |
|
|
| logger = get_logger(__name__) |
|
|
|
|
| @attr.s(auto_attribs=True) |
| class ModelCheckpoint: |
| steps: int |
| file_path: str |
| reward: Optional[float] |
| creation_time: float |
| auxillary_file_paths: List[str] = attr.ib(factory=list) |
|
|
|
|
| class ModelCheckpointManager: |
| @staticmethod |
| def get_checkpoints(behavior_name: str) -> List[Dict[str, Any]]: |
| checkpoint_list = GlobalTrainingStatus.get_parameter_state( |
| behavior_name, StatusType.CHECKPOINTS |
| ) |
| if not checkpoint_list: |
| checkpoint_list = [] |
| GlobalTrainingStatus.set_parameter_state( |
| behavior_name, StatusType.CHECKPOINTS, checkpoint_list |
| ) |
| return checkpoint_list |
|
|
| @staticmethod |
| def remove_checkpoint(checkpoint: Dict[str, Any]) -> None: |
| """ |
| Removes a checkpoint stored in checkpoint_list. |
| If checkpoint cannot be found, no action is done. |
| |
| :param checkpoint: A checkpoint stored in checkpoint_list |
| """ |
| file_paths: List[str] = [checkpoint["file_path"]] |
| file_paths.extend(checkpoint["auxillary_file_paths"]) |
| for file_path in file_paths: |
| if os.path.exists(file_path): |
| os.remove(file_path) |
| logger.debug(f"Removed checkpoint model {file_path}.") |
| else: |
| logger.debug(f"Checkpoint at {file_path} could not be found.") |
| return |
|
|
| @classmethod |
| def _cleanup_extra_checkpoints( |
| cls, checkpoints: List[Dict], keep_checkpoints: int |
| ) -> List[Dict]: |
| """ |
| Ensures that the number of checkpoints stored are within the number |
| of checkpoints the user defines. If the limit is hit, checkpoints are |
| removed to create room for the next checkpoint to be inserted. |
| |
| :param behavior_name: The behavior name whose checkpoints we will mange. |
| :param keep_checkpoints: Number of checkpoints to record (user-defined). |
| """ |
| while len(checkpoints) > keep_checkpoints: |
| if keep_checkpoints <= 0 or len(checkpoints) == 0: |
| break |
| ModelCheckpointManager.remove_checkpoint(checkpoints.pop(0)) |
| return checkpoints |
|
|
| @classmethod |
| def add_checkpoint( |
| cls, behavior_name: str, new_checkpoint: ModelCheckpoint, keep_checkpoints: int |
| ) -> None: |
| """ |
| Make room for new checkpoint if needed and insert new checkpoint information. |
| :param behavior_name: Behavior name for the checkpoint. |
| :param new_checkpoint: The new checkpoint to be recorded. |
| :param keep_checkpoints: Number of checkpoints to record (user-defined). |
| """ |
| new_checkpoint_dict = attr.asdict(new_checkpoint) |
| checkpoints = cls.get_checkpoints(behavior_name) |
| checkpoints.append(new_checkpoint_dict) |
| cls._cleanup_extra_checkpoints(checkpoints, keep_checkpoints) |
| GlobalTrainingStatus.set_parameter_state( |
| behavior_name, StatusType.CHECKPOINTS, checkpoints |
| ) |
|
|
| @classmethod |
| def track_final_checkpoint( |
| cls, behavior_name: str, final_checkpoint: ModelCheckpoint |
| ) -> None: |
| """ |
| Ensures number of checkpoints stored is within the max number of checkpoints |
| defined by the user and finally stores the information about the final |
| model (or intermediate model if training is interrupted). |
| :param behavior_name: Behavior name of the model. |
| :param final_checkpoint: Checkpoint information for the final model. |
| """ |
| final_model_dict = attr.asdict(final_checkpoint) |
| GlobalTrainingStatus.set_parameter_state( |
| behavior_name, StatusType.FINAL_CHECKPOINT, final_model_dict |
| ) |
|
|
