(rank,
load_save_dict: dict = None,
model_for_opt: tuple = None,
optimizer=None,
ckpt_dir: str = None,
load_optimizer_only=False,
lora=False)
| 209 | |
| 210 | @staticmethod |
| 211 | def load_ckpt_from_dir(rank, |
| 212 | load_save_dict: dict = None, |
| 213 | model_for_opt: tuple = None, |
| 214 | optimizer=None, |
| 215 | ckpt_dir: str = None, |
| 216 | load_optimizer_only=False, |
| 217 | lora=False): |
| 218 | if not load_optimizer_only: |
| 219 | if load_save_dict is not None: |
| 220 | for model_name, model in load_save_dict.items(): |
| 221 | file_path = os.path.join(ckpt_dir, f'{model_name}.pt') |
| 222 | folder_path = os.path.join(ckpt_dir, model_name) |
| 223 | load_path = file_path if os.path.isfile( |
| 224 | file_path) else folder_path |
| 225 | TrainerBase.load_model( |
| 226 | rank, model, load_path, lora, strict=False, logger=None) |
| 227 | if optimizer is not None: |
| 228 | for model_name in model_for_opt: |
| 229 | file_path = os.path.join(ckpt_dir, |
| 230 | f'{model_name}-opt-{rank:02d}.pt') |
| 231 | folder_path = os.path.join(ckpt_dir, f'{model_name}-opt') |
| 232 | load_path = file_path if os.path.isfile( |
| 233 | file_path) else folder_path |
| 234 | TrainerBase.load_optimizer(load_save_dict[model_name], |
| 235 | optimizer, load_path) |
| 236 | return None |
| 237 | |
| 238 | def load_ckpt(self, |
| 239 | rank, |
no test coverage detected