mirror of
https://github.com/AIGC-Audio/AudioGPT.git
synced 2026-07-10 04:20:09 +02:00
69 lines
2.7 KiB
Python
Executable File
69 lines
2.7 KiB
Python
Executable File
import glob
|
|
import logging
|
|
import os
|
|
import re
|
|
import torch
|
|
|
|
|
|
def get_last_checkpoint(work_dir, steps=None):
|
|
checkpoint = None
|
|
last_ckpt_path = None
|
|
ckpt_paths = get_all_ckpts(work_dir, steps)
|
|
if len(ckpt_paths) > 0:
|
|
last_ckpt_path = ckpt_paths[0]
|
|
checkpoint = torch.load(last_ckpt_path, map_location='cpu')
|
|
logging.info(f'load module from checkpoint: {last_ckpt_path}')
|
|
return checkpoint, last_ckpt_path
|
|
|
|
|
|
def get_all_ckpts(work_dir, steps=None):
|
|
if steps is None:
|
|
ckpt_path_pattern = f'{work_dir}/model_ckpt_steps_*.ckpt'
|
|
else:
|
|
ckpt_path_pattern = f'{work_dir}/model_ckpt_steps_{steps}.ckpt'
|
|
return sorted(glob.glob(ckpt_path_pattern),
|
|
key=lambda x: -int(re.findall('.*steps\_(\d+)\.ckpt', x)[0]))
|
|
|
|
|
|
def load_ckpt(cur_model, ckpt_base_dir, model_name='model', force=True, strict=True):
|
|
if os.path.isfile(ckpt_base_dir):
|
|
base_dir = os.path.dirname(ckpt_base_dir)
|
|
ckpt_path = ckpt_base_dir
|
|
checkpoint = torch.load(ckpt_base_dir, map_location='cpu')
|
|
else:
|
|
base_dir = ckpt_base_dir
|
|
checkpoint, ckpt_path = get_last_checkpoint(ckpt_base_dir)
|
|
if checkpoint is not None:
|
|
state_dict = checkpoint["state_dict"]
|
|
if len([k for k in state_dict.keys() if '.' in k]) > 0:
|
|
state_dict = {k[len(model_name) + 1:]: v for k, v in state_dict.items()
|
|
if k.startswith(f'{model_name}.')}
|
|
else:
|
|
if '.' not in model_name:
|
|
state_dict = state_dict[model_name]
|
|
else:
|
|
base_model_name = model_name.split('.')[0]
|
|
rest_model_name = model_name[len(base_model_name) + 1:]
|
|
state_dict = {
|
|
k[len(rest_model_name) + 1:]: v for k, v in state_dict[base_model_name].items()
|
|
if k.startswith(f'{rest_model_name}.')}
|
|
if not strict:
|
|
cur_model_state_dict = cur_model.state_dict()
|
|
unmatched_keys = []
|
|
for key, param in state_dict.items():
|
|
if key in cur_model_state_dict:
|
|
new_param = cur_model_state_dict[key]
|
|
if new_param.shape != param.shape:
|
|
unmatched_keys.append(key)
|
|
print("| Unmatched keys: ", key, new_param.shape, param.shape)
|
|
for key in unmatched_keys:
|
|
del state_dict[key]
|
|
cur_model.load_state_dict(state_dict, strict=strict)
|
|
print(f"| load '{model_name}' from '{ckpt_path}'.")
|
|
else:
|
|
e_msg = f"| ckpt not found in {base_dir}."
|
|
if force:
|
|
assert False, e_msg
|
|
else:
|
|
print(e_msg)
|