bytetrack / yolox /utils /checkpoint.py
AK391
all files
7734d5b
raw
history blame
1.33 kB
#!/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)