text/init.py (#203)
* Create all_process.py * Create asr_transcript.py * Update config.py * Create extract_list.py * Create clean_list.py * Create custom.css * Create compress_model.py * Update all_process.py * Update resample.py * Update resample.py * configs/config.json copy utils * mirror: openi + token bert models optimize * text/__init__.py platform compatibility compress_model.py output * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
88
compress_model.py
Normal file
88
compress_model.py
Normal file
@@ -0,0 +1,88 @@
|
||||
from collections import OrderedDict
|
||||
from text.symbols import symbols
|
||||
import torch
|
||||
from tools.log import logger
|
||||
import utils
|
||||
from models import SynthesizerTrn
|
||||
import os
|
||||
|
||||
|
||||
def copyStateDict(state_dict):
|
||||
if list(state_dict.keys())[0].startswith("module"):
|
||||
start_idx = 1
|
||||
else:
|
||||
start_idx = 0
|
||||
new_state_dict = OrderedDict()
|
||||
for k, v in state_dict.items():
|
||||
name = ",".join(k.split(".")[start_idx:])
|
||||
new_state_dict[name] = v
|
||||
return new_state_dict
|
||||
|
||||
|
||||
def removeOptimizer(config: str, input_model: str, ishalf: bool, output_model: str):
|
||||
hps = utils.get_hparams_from_file(config)
|
||||
|
||||
net_g = SynthesizerTrn(
|
||||
len(symbols),
|
||||
hps.data.filter_length // 2 + 1,
|
||||
hps.train.segment_size // hps.data.hop_length,
|
||||
n_speakers=hps.data.n_speakers,
|
||||
**hps.model,
|
||||
)
|
||||
|
||||
optim_g = torch.optim.AdamW(
|
||||
net_g.parameters(),
|
||||
hps.train.learning_rate,
|
||||
betas=hps.train.betas,
|
||||
eps=hps.train.eps,
|
||||
)
|
||||
|
||||
state_dict_g = torch.load(input_model, map_location="cpu")
|
||||
new_dict_g = copyStateDict(state_dict_g)
|
||||
keys = []
|
||||
for k, v in new_dict_g["model"].items():
|
||||
if "enc_q" in k:
|
||||
continue # noqa: E701
|
||||
keys.append(k)
|
||||
|
||||
new_dict_g = (
|
||||
{k: new_dict_g["model"][k].half() for k in keys}
|
||||
if ishalf
|
||||
else {k: new_dict_g["model"][k] for k in keys}
|
||||
)
|
||||
|
||||
torch.save(
|
||||
{
|
||||
"model": new_dict_g,
|
||||
"iteration": 0,
|
||||
"optimizer": optim_g.state_dict(),
|
||||
"learning_rate": 0.0001,
|
||||
},
|
||||
output_model,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("-c", "--config", type=str, default="configs/config.json")
|
||||
parser.add_argument("-i", "--input", type=str)
|
||||
parser.add_argument("-o", "--output", type=str, default=None)
|
||||
parser.add_argument(
|
||||
"-hf", "--half", action="store_true", default=False, help="Save as FP16"
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
output = args.output
|
||||
|
||||
if output is None:
|
||||
import os.path
|
||||
|
||||
filename, ext = os.path.splitext(args.input)
|
||||
half = "_half" if args.half else ""
|
||||
output = filename + "_release" + half + ext
|
||||
|
||||
removeOptimizer(args.config, args.input, args.half, output)
|
||||
logger.info(f"压缩模型成功, 输出模型: {os.path.abspath(output)}")
|
||||
Reference in New Issue
Block a user