| 10 | |
| 11 | |
| 12 | def load_ckpt(model, ckpt): |
| 13 | model_state_dict = model.state_dict() |
| 14 | load_dict = {} |
| 15 | for key_model, v in model_state_dict.items(): |
| 16 | if key_model not in ckpt: |
| 17 | logger.warning( |
| 18 | "{} is not in the ckpt. Please double check and see if this is desired.".format( |
| 19 | key_model |
| 20 | ) |
| 21 | ) |
| 22 | continue |
| 23 | v_ckpt = ckpt[key_model] |
| 24 | if v.shape != v_ckpt.shape: |
| 25 | logger.warning( |
| 26 | "Shape of {} in checkpoint is {}, while shape of {} in model is {}.".format( |
| 27 | key_model, v_ckpt.shape, key_model, v.shape |
| 28 | ) |
| 29 | ) |
| 30 | continue |
| 31 | load_dict[key_model] = v_ckpt |
| 32 | |
| 33 | model.load_state_dict(load_dict, strict=False) |
| 34 | return model |
| 35 | |
| 36 | |
| 37 | def save_checkpoint(state, is_best, save_dir, model_name=""): |