89 lines
3.3 KiB
Python
89 lines
3.3 KiB
Python
|
import os
|
|||
|
|
import re
|
|||
|
|
import shutil
|
|||
|
|
|
|||
|
|
import torch
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_model_folder(path):
|
|||
|
|
model_lists = os.listdir(path)
|
|||
|
|
res_list = []
|
|||
|
|
filter_list = ["hubert", "xiaoma_pe", "hifigan", "checkpoints", ".yaml", ".zip"]
|
|||
|
|
for path in model_lists:
|
|||
|
|
if not any(word if word in path else False for word in filter_list):
|
|||
|
|
res_list.append(path)
|
|||
|
|
return res_list
|
|||
|
|
|
|||
|
|
|
|||
|
|
def scan(path):
|
|||
|
|
model_str = ""
|
|||
|
|
path_lists = get_model_folder(path)
|
|||
|
|
for i in range(0, len(path_lists)):
|
|||
|
|
if re.search(u'[\u4e00-\u9fa5]', path_lists[i]):
|
|||
|
|
print(f'{path_lists[i]}:中文路径！此项跳过')
|
|||
|
|
continue
|
|||
|
|
model_str += f"{i}:{path_lists[i]} "
|
|||
|
|
if (i + 1) % 5 == 0:
|
|||
|
|
print(f"{model_str}")
|
|||
|
|
model_str = ""
|
|||
|
|
if len(path_lists) % 5 != 0:
|
|||
|
|
print(model_str)
|
|||
|
|
return path_lists
|
|||
|
|
|
|||
|
|
|
|||
|
|
def simplify_pth(model_name, proj_name, output_path):
|
|||
|
|
model_path = f'./checkpoints/{proj_name}'
|
|||
|
|
checkpoint_dict = torch.load(f'{model_path}/{model_name}')
|
|||
|
|
torch.save({'epoch': checkpoint_dict['epoch'],
|
|||
|
|
'state_dict': checkpoint_dict['state_dict'],
|
|||
|
|
'global_step': None,
|
|||
|
|
'checkpoint_callback_best': None,
|
|||
|
|
'optimizer_states': None,
|
|||
|
|
'lr_schedulers': None
|
|||
|
|
}, output_path)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def mkdir(paths: list):
|
|||
|
|
for path in paths:
|
|||
|
|
if not os.path.exists(path):
|
|||
|
|
os.mkdir(path)
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == '__main__':
|
|||
|
|
if os.path.exists("./checkpoints"):
|
|||
|
|
path_list = scan("./checkpoints")
|
|||
|
|
else:
|
|||
|
|
print("请检查checkpoints文件夹是否存在")
|
|||
|
|
exit()
|
|||
|
|
a = input("\r\n请输入序号并回车:")
|
|||
|
|
project_name = path_list[int(a)]
|
|||
|
|
path_list = scan(f"./checkpoints/{path_list[int(a)]}")
|
|||
|
|
b = input("\r\n请输入序号并回车:")
|
|||
|
|
pth_name = path_list[int(b)]
|
|||
|
|
|
|||
|
|
print("\r\n选择:\r\n"
|
|||
|
|
"0.存储精简模型到对应模型目录（本地精简模型时推荐使用这个）\r\n"
|
|||
|
|
"1.存储精简模型和config.yaml到程序根目录（新建文件夹，九天毕昇上导出精简模型推荐使用这个）\r\n"
|
|||
|
|
"2.复制完整模型和config.yaml到程序根目录（新建文件夹，九天毕昇上导出完整模型推荐使用这个）\r\n"
|
|||
|
|
"输入其他退出")
|
|||
|
|
f = int(input("\r\n请输入序号并回车:"))
|
|||
|
|
if f == 0:
|
|||
|
|
print(f"已保存精简模型至对应模型目录")
|
|||
|
|
shutil.copyfile(f'./checkpoints/{project_name}/config.yaml', f"./{project_name}/config.yaml")
|
|||
|
|
output = f"./checkpoints/{project_name}/clean_{pth_name}"
|
|||
|
|
simplify_pth(pth_name, project_name, output)
|
|||
|
|
elif f == 1:
|
|||
|
|
print(f"已保存精简模型至: 根目录下新建文件夹/{project_name}")
|
|||
|
|
mkdir([f"./{project_name}"])
|
|||
|
|
shutil.copyfile(f'./checkpoints/{project_name}/config.yaml', f"./{project_name}/config.yaml")
|
|||
|
|
output = f"./{project_name}/clean_{pth_name}"
|
|||
|
|
simplify_pth(pth_name, project_name, output)
|
|||
|
|
elif f == 2:
|
|||
|
|
print(f"已保存完整模型至: 根目录下新建文件夹/{project_name}")
|
|||
|
|
mkdir([f"./{project_name}"])
|
|||
|
|
shutil.copyfile(f'./checkpoints/{project_name}/config.yaml', f"./{project_name}/config.yaml")
|
|||
|
|
shutil.copyfile(f'./checkpoints/{project_name}/{pth_name}', f"./{project_name}/{pth_name}")
|
|||
|
|
else:
|
|||
|
|
print("输入错误，程序退出")
|
|||
|
|
exit()
|