修改加载latest model的方式,修改global_step计算,增加preserved参数,增加train_with_pretrained_model参数

This commit is contained in:
Emberstar
2023-06-14 10:18:19 +08:00
parent e97b185188
commit 576424fe58
2 changed files with 46 additions and 14 deletions
+24 -3
View File
@@ -204,14 +204,29 @@ def summarize(writer, global_step, scalars={}, histograms={}, images={}, audios=
writer.add_audio(k, v, global_step, audio_sampling_rate)
def latest_checkpoint_path(dir_path, regex="G_*.pth"):
def extract_digits(f):
digits = "".join(filter(str.isdigit, f))
return int(digits) if digits else -1
def latest_checkpoint_path(dir_path, regex="G_[0-9]*.pth"):
f_list = glob.glob(os.path.join(dir_path, regex))
f_list.sort(key=lambda f: int("".join(filter(str.isdigit, f))))
f_list.sort(key=lambda f: extract_digits(f))
x = f_list[-1]
print(x)
print(f"latest_checkpoint_path:{x}")
return x
def oldest_checkpoint_path(dir_path, regex="G_[0-9]*.pth", preserved=4):
f_list = glob.glob(os.path.join(dir_path, regex))
f_list.sort(key=lambda f: extract_digits(f))
if len(f_list) > preserved:
x = f_list[0]
print(f"oldest_checkpoint_path:{x}")
return x
return ""
def plot_spectrogram_to_numpy(spectrogram):
global MATPLOTLIB_FLAG
if not MATPLOTLIB_FLAG:
@@ -288,6 +303,10 @@ def get_hparams(init=True):
help='finetune epochs')
parser.add_argument('--cont', type=bool, default=False, help='whether to continue training on the latest checkpoint')
parser.add_argument('--drop_speaker_embed', type=bool, default=False, help='whether to drop existing characters')
parser.add_argument('--train_with_pretrained_model', type=bool, default=True,
help='whether to train with pretrained model')
parser.add_argument('--preserved', type=int, default=4,
help='Number of preserved models')
args = parser.parse_args()
model_dir = os.path.join("./", args.model)
@@ -312,6 +331,8 @@ def get_hparams(init=True):
hparams.max_epochs = args.max_epochs
hparams.cont = args.cont
hparams.drop_speaker_embed = args.drop_speaker_embed
hparams.train_with_pretrained_model = args.train_with_pretrained_model
hparams.preserved = args.preserved
return hparams