Spaces:
Runtime error
Runtime error
File size: 1,331 Bytes
7734d5b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
from loguru import logger
import torch
import os
import shutil
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.tar")
torch.save(state, filename)
if is_best:
best_filename = os.path.join(save_dir, "best_ckpt.pth.tar")
shutil.copyfile(filename, best_filename)
|