|
|
|
|
|
|
|
import os |
|
import shutil |
|
from loguru import logger |
|
|
|
import torch |
|
|
|
|
|
def load_ckpt(model, ckpt): |
|
model_state_dict = model.state_dict() |
|
load_dict = {} |
|
for key_model, v in model_state_dict.items(): |
|
if key_model not in ckpt: |
|
logger.warning( |
|
"{} is not in the ckpt. Please double check and see if this is desired.".format( |
|
key_model |
|
) |
|
) |
|
continue |
|
v_ckpt = ckpt[key_model] |
|
if v.shape != v_ckpt.shape: |
|
logger.warning( |
|
"Shape of {} in checkpoint is {}, while shape of {} in model is {}.".format( |
|
key_model, v_ckpt.shape, key_model, v.shape |
|
) |
|
) |
|
continue |
|
load_dict[key_model] = v_ckpt |
|
|
|
model.load_state_dict(load_dict, strict=False) |
|
return model |
|
|
|
|
|
def save_checkpoint(state, is_best, save_dir, model_name=""): |
|
if not os.path.exists(save_dir): |
|
os.makedirs(save_dir) |
|
filename = os.path.join(save_dir, model_name + "_ckpt.pth") |
|
torch.save(state, filename) |
|
if is_best: |
|
best_filename = os.path.join(save_dir, "best_ckpt.pth") |
|
shutil.copyfile(filename, best_filename) |
|
|
