Spaces:
Sleeping
Sleeping
| import torch | |
| import copy | |
| def resume_and_load(model, ckpt_path, device): | |
| print("Loading checkpoints from", ckpt_path) | |
| checkpoints = torch.load(ckpt_path, map_location=device) | |
| if 'model' in checkpoints.keys() and 'optimizer' in checkpoints.keys(): | |
| checkpoints = convert_official_ckpt(checkpoints, model.state_dict()) | |
| missing_keys, unexpected_keys = model.load_state_dict(checkpoints) | |
| print("Missing keys:", missing_keys) | |
| print("Unexpected keys:", unexpected_keys) | |
| return model | |
| def save_ckpt(model, save_path, distributed=False): | |
| print("Saving checkpoints to", save_path) | |
| state_dict = model.state_dict() if not distributed else model.module.state_dict() | |
| for k in list(state_dict.keys()): | |
| if "domain" in k or "mae" in k: | |
| state_dict.pop(k) | |
| torch.save(state_dict, save_path) | |
| def selective_reinitialize(model, reinit_ckpt, keep_modules): | |
| print("Doing selective reinitialization. Parameters of the model will be reinitialized EXCEPT FOR:") | |
| for key in copy.deepcopy(list(reinit_ckpt.keys())): | |
| to_be_reinit = True | |
| for keep_module in keep_modules: | |
| if keep_module in key: | |
| to_be_reinit = False | |
| break | |
| if not to_be_reinit: | |
| reinit_ckpt.pop(key) | |
| print(key) | |
| model.load_state_dict(reinit_ckpt, strict=False) | |
| return model | |
| def convert_official_ckpt(checkpoints, state_dict): | |
| checkpoints = checkpoints['model'] | |
| official_keys, new_keys = sorted(list(checkpoints.keys())), sorted(list(state_dict.keys())) | |
| new_state_dict = {} | |
| for k_official, k_new in zip(official_keys, new_keys): | |
| if not k_official.startswith('class'): | |
| new_state_dict[k_new] = checkpoints[k_official] | |
| else: | |
| print("Skipping", k_official) | |
| new_state_dict[k_new] = state_dict[k_new] | |
| return new_state_dict | |