File size: 1,799 Bytes
747451d | 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 46 47 48 49 50 51 52 | # /*---------------------------------------------------------------------------------------------
# * Copyright (c) 2025 STMicroelectronics.
# * All rights reserved.
# *
# * This software is licensed under terms that can be found in the LICENSE file in
# * the root directory of this software component.
# * If no LICENSE file comes with this software, it is provided AS-IS.
# *--------------------------------------------------------------------------------------------*/
#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# Copyright (c) Megvii Inc. All rights reserved.
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)
|