| from typing import Optional, Dict |
| import os |
|
|
| class TopKCheckpointManager: |
| def __init__(self, |
| save_dir, |
| monitor_key: str, |
| mode='min', |
| k=1, |
| format_str='epoch={epoch:03d}-train_loss={train_loss:.3f}.ckpt' |
| ): |
| assert mode in ['max', 'min'] |
| assert k >= 0 |
|
|
| self.save_dir = save_dir |
| self.monitor_key = monitor_key |
| self.mode = mode |
| self.k = k |
| self.format_str = format_str |
| self.path_value_map = dict() |
| |
| def get_ckpt_path(self, data: Dict[str, float]) -> Optional[str]: |
| if self.k == 0: |
| return None |
|
|
| value = data[self.monitor_key] |
| ckpt_path = os.path.join( |
| self.save_dir, self.format_str.format(**data)) |
| |
| if len(self.path_value_map) < self.k: |
| |
| self.path_value_map[ckpt_path] = value |
| return ckpt_path |
| |
| |
| sorted_map = sorted(self.path_value_map.items(), key=lambda x: x[1]) |
| min_path, min_value = sorted_map[0] |
| max_path, max_value = sorted_map[-1] |
|
|
| delete_path = None |
| if self.mode == 'max': |
| if value > min_value: |
| delete_path = min_path |
| else: |
| if value < max_value: |
| delete_path = max_path |
|
|
| if delete_path is None: |
| return None |
| else: |
| del self.path_value_map[delete_path] |
| self.path_value_map[ckpt_path] = value |
|
|
| if not os.path.exists(self.save_dir): |
| os.mkdir(self.save_dir) |
|
|
| if os.path.exists(delete_path): |
| os.remove(delete_path) |
| return ckpt_path |
|
|
