Paolify/RVC_4
0
1import os, traceback2import glob3import sys4import argparse5import logging6import json7import subprocess8import numpy as np9from scipy.io.wavfile import read10import torch11 12MATPLOTLIB_FLAG = False13 14logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)15logger = logging16 17 18def load_checkpoint_d(checkpoint_path, combd, sbd, optimizer=None, load_opt=1):19 assert os.path.isfile(checkpoint_path)20 checkpoint_dict = torch.load(checkpoint_path, map_location="cpu")21 22 ##################23 def go(model, bkey):24 saved_state_dict = checkpoint_dict[bkey]25 if hasattr(model, "module"):26 state_dict = model.module.state_dict()27 else:28 state_dict = model.state_dict()29 new_state_dict = {}30 for k, v in state_dict.items(): # 模型需要的shape31 try:32 new_state_dict[k] = saved_state_dict[k]33 if saved_state_dict[k].shape != state_dict[k].shape:34 print(35 "shape-%s-mismatch|need-%s|get-%s"36 % (k, state_dict[k].shape, saved_state_dict[k].shape)37 ) #38 raise KeyError39 except:40 # logger.info(traceback.format_exc())41 logger.info("%s is not in the checkpoint" % k) # pretrain缺失的42 new_state_dict[k] = v # 模型自带的随机值43 if hasattr(model, "module"):44 model.module.load_state_dict(new_state_dict, strict=False)45 else:46 model.load_state_dict(new_state_dict, strict=False)47 48 go(combd, "combd")49 go(sbd, "sbd")50 #############51 logger.info("Loaded model weights")52 53 iteration = checkpoint_dict["iteration"]54 learning_rate = checkpoint_dict["learning_rate"]55 if (56 optimizer is not None and load_opt == 157 ): ###加载不了，如果是空的的话，重新初始化，可能还会影响lr时间表的更新，因此在train文件最外围catch58 # try:59 optimizer.load_state_dict(checkpoint_dict["optimizer"])60 # except:61 # traceback.print_exc()62 logger.info("Loaded checkpoint '{}' (epoch {})".format(checkpoint_path, iteration))63 return model, optimizer, learning_rate, iteration64 65 66# def load_checkpoint(checkpoint_path, model, optimizer=None):67# assert os.path.isfile(checkpoint_path)68# checkpoint_dict = torch.load(checkpoint_path, map_location='cpu')69# iteration = checkpoint_dict['iteration']70# learning_rate = checkpoint_dict['learning_rate']71# if optimizer is not None:72# optimizer.load_state_dict(checkpoint_dict['optimizer'])73# # print(1111)74# saved_state_dict = checkpoint_dict['model']75# # print(1111)76#77# if hasattr(model, 'module'):78# state_dict = model.module.state_dict()79# else:80# state_dict = model.state_dict()81# new_state_dict= {}82# for k, v in state_dict.items():83# try:84# new_state_dict[k] = saved_state_dict[k]85# except:86# logger.info("%s is not in the checkpoint" % k)87# new_state_dict[k] = v88# if hasattr(model, 'module'):89# model.module.load_state_dict(new_state_dict)90# else:91# model.load_state_dict(new_state_dict)92# logger.info("Loaded checkpoint '{}' (epoch {})" .format(93# checkpoint_path, iteration))94# return model, optimizer, learning_rate, iteration95def load_checkpoint(checkpoint_path, model, optimizer=None, load_opt=1):96 assert os.path.isfile(checkpoint_path)97 checkpoint_dict = torch.load(checkpoint_path, map_location="cpu")98 99 saved_state_dict = checkpoint_dict["model"]100 if hasattr(model, "module"):101 state_dict = model.module.state_dict()102 else:103 state_dict = model.state_dict()104 new_state_dict = {}105 for k, v in state_dict.items(): # 模型需要的shape106 try:107 new_state_dict[k] = saved_state_dict[k]108 if saved_state_dict[k].shape != state_dict[k].shape:109 print(110 "shape-%s-mismatch|need-%s|get-%s"111 % (k, state_dict[k].shape, saved_state_dict[k].shape)112 ) #113 raise KeyError114 except:115 # logger.info(traceback.format_exc())116 logger.info("%s is not in the checkpoint" % k) # pretrain缺失的117 new_state_dict[k] = v # 模型自带的随机值118 if hasattr(model, "module"):119 model.module.load_state_dict(new_state_dict, strict=False)120 else:121 model.load_state_dict(new_state_dict, strict=False)122 logger.info("Loaded model weights")123 124 iteration = checkpoint_dict["iteration"]125 learning_rate = checkpoint_dict["learning_rate"]126 if (127 optimizer is not None and load_opt == 1128 ): ###加载不了，如果是空的的话，重新初始化，可能还会影响lr时间表的更新，因此在train文件最外围catch129 # try:130 optimizer.load_state_dict(checkpoint_dict["optimizer"])131 # except:132 # traceback.print_exc()133 logger.info("Loaded checkpoint '{}' (epoch {})".format(checkpoint_path, iteration))134 return model, optimizer, learning_rate, iteration135 136 137def save_checkpoint(model, optimizer, learning_rate, iteration, checkpoint_path):138 logger.info(139 "Saving model and optimizer state at epoch {} to {}".format(140 iteration, checkpoint_path141 )142 )143 if hasattr(model, "module"):144 state_dict = model.module.state_dict()145 else:146 state_dict = model.state_dict()147 torch.save(148 {149 "model": state_dict,150 "iteration": iteration,151 "optimizer": optimizer.state_dict(),152 "learning_rate": learning_rate,153 },154 checkpoint_path,155 )156 157 158def save_checkpoint_d(combd, sbd, optimizer, learning_rate, iteration, checkpoint_path):159 logger.info(160 "Saving model and optimizer state at epoch {} to {}".format(161 iteration, checkpoint_path162 )163 )164 if hasattr(combd, "module"):165 state_dict_combd = combd.module.state_dict()166 else:167 state_dict_combd = combd.state_dict()168 if hasattr(sbd, "module"):169 state_dict_sbd = sbd.module.state_dict()170 else:171 state_dict_sbd = sbd.state_dict()172 torch.save(173 {174 "combd": state_dict_combd,175 "sbd": state_dict_sbd,176 "iteration": iteration,177 "optimizer": optimizer.state_dict(),178 "learning_rate": learning_rate,179 },180 checkpoint_path,181 )182 183 184def summarize(185 writer,186 global_step,187 scalars={},188 histograms={},189 images={},190 audios={},191 audio_sampling_rate=22050,192):193 for k, v in scalars.items():194 writer.add_scalar(k, v, global_step)195 for k, v in histograms.items():196 writer.add_histogram(k, v, global_step)197 for k, v in images.items():198 writer.add_image(k, v, global_step, dataformats="HWC")199 for k, v in audios.items():200 writer.add_audio(k, v, global_step, audio_sampling_rate)201 202 203def latest_checkpoint_path(dir_path, regex="G_*.pth"):204 f_list = glob.glob(os.path.join(dir_path, regex))205 f_list.sort(key=lambda f: int("".join(filter(str.isdigit, f))))206 x = f_list[-1]207 print(x)208 return x209 210 211def plot_spectrogram_to_numpy(spectrogram):212 global MATPLOTLIB_FLAG213 if not MATPLOTLIB_FLAG:214 import matplotlib215 216 matplotlib.use("Agg")217 MATPLOTLIB_FLAG = True218 mpl_logger = logging.getLogger("matplotlib")219 mpl_logger.setLevel(logging.WARNING)220 import matplotlib.pylab as plt221 import numpy as np222 223 fig, ax = plt.subplots(figsize=(10, 2))224 im = ax.imshow(spectrogram, aspect="auto", origin="lower", interpolation="none")225 plt.colorbar(im, ax=ax)226 plt.xlabel("Frames")227 plt.ylabel("Channels")228 plt.tight_layout()229 230 fig.canvas.draw()231 data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep="")232 data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))233 plt.close()234 return data235 236 237def plot_alignment_to_numpy(alignment, info=None):238 global MATPLOTLIB_FLAG239 if not MATPLOTLIB_FLAG:240 import matplotlib241 242 matplotlib.use("Agg")243 MATPLOTLIB_FLAG = True244 mpl_logger = logging.getLogger("matplotlib")245 mpl_logger.setLevel(logging.WARNING)246 import matplotlib.pylab as plt247 import numpy as np248 249 fig, ax = plt.subplots(figsize=(6, 4))250 im = ax.imshow(251 alignment.transpose(), aspect="auto", origin="lower", interpolation="none"252 )253 fig.colorbar(im, ax=ax)254 xlabel = "Decoder timestep"255 if info is not None:256 xlabel += "\n\n" + info257 plt.xlabel(xlabel)258 plt.ylabel("Encoder timestep")259 plt.tight_layout()260 261 fig.canvas.draw()262 data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep="")263 data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))264 plt.close()265 return data266 267 268def load_wav_to_torch(full_path):269 sampling_rate, data = read(full_path)270 return torch.FloatTensor(data.astype(np.float32)), sampling_rate271 272 273def load_filepaths_and_text(filename, split="|"):274 with open(filename, encoding='utf-8') as f:275 filepaths_and_text = [line.strip().split(split) for line in f]276 filepaths_and_text = [item for item in filepaths_and_text if len(item) == 5] # ensure there are 5 items.277 return filepaths_and_text278 279 280def get_hparams(init=True):281 """282 todo:283 结尾七人组：284 保存频率、总epoch done285 bs done286 pretrainG、pretrainD done287 卡号：os.en["CUDA_VISIBLE_DEVICES"] done288 if_latest done289 模型：if_f0 done290 采样率：自动选择config done291 是否缓存数据集进GPU:if_cache_data_in_gpu done292 293 -m:294 自动决定training_files路径,改掉train_nsf_load_pretrain.py里的hps.data.training_files done295 -c不要了296 """297 parser = argparse.ArgumentParser()298 # parser.add_argument('-c', '--config', type=str, default="configs/40k.json",help='JSON file for configuration')299 parser.add_argument(300 "-se",301 "--save_every_epoch",302 type=int,303 required=True,304 help="checkpoint save frequency (epoch)",305 )306 parser.add_argument(307 "-te", "--total_epoch", type=int, required=True, help="total_epoch"308 )309 parser.add_argument(310 "-pg", "--pretrainG", type=str, default="", help="Pretrained Discriminator path"311 )312 parser.add_argument(313 "-pd", "--pretrainD", type=str, default="", help="Pretrained Generator path"314 )315 parser.add_argument("-g", "--gpus", type=str, default="0", help="split by -")316 parser.add_argument(317 "-bs", "--batch_size", type=int, required=True, help="batch size"318 )319 parser.add_argument(320 "-e", "--experiment_dir", type=str, required=True, help="experiment dir"321 ) # -m322 parser.add_argument(323 "-sr", "--sample_rate", type=str, required=True, help="sample rate, 32k/40k/48k"324 )325 parser.add_argument(326 "-sw",327 "--save_every_weights",328 type=str,329 default="0",330 help="save the extracted model in weights directory when saving checkpoints",331 )332 parser.add_argument(333 "-v", "--version", type=str, required=True, help="model version"334 )335 parser.add_argument(336 "-f0",337 "--if_f0",338 type=int,339 required=True,340 help="use f0 as one of the inputs of the model, 1 or 0",341 )342 parser.add_argument(343 "-l",344 "--if_latest",345 type=int,346 required=True,347 help="if only save the latest G/D pth file, 1 or 0",348 )349 parser.add_argument(350 "-c",351 "--if_cache_data_in_gpu",352 type=int,353 required=True,354 help="if caching the dataset in GPU memory, 1 or 0",355 )356 parser.add_argument(357 "-li", "--log_interval", type=int, required=True, help="log interval"358 )359 360 args = parser.parse_args()361 name = args.experiment_dir362 experiment_dir = os.path.join("./logs", args.experiment_dir)363 364 if not os.path.exists(experiment_dir):365 os.makedirs(experiment_dir)366 367 if args.version == "v1" or args.sample_rate == "40k":368 config_path = "configs/%s.json" % args.sample_rate369 else:370 config_path = "configs/%s_v2.json" % args.sample_rate371 config_save_path = os.path.join(experiment_dir, "config.json")372 if init:373 with open(config_path, "r") as f:374 data = f.read()375 with open(config_save_path, "w") as f:376 f.write(data)377 else:378 with open(config_save_path, "r") as f:379 data = f.read()380 config = json.loads(data)381 382 hparams = HParams(**config)383 hparams.model_dir = hparams.experiment_dir = experiment_dir384 hparams.save_every_epoch = args.save_every_epoch385 hparams.name = name386 hparams.total_epoch = args.total_epoch387 hparams.pretrainG = args.pretrainG388 hparams.pretrainD = args.pretrainD389 hparams.version = args.version390 hparams.gpus = args.gpus391 hparams.train.batch_size = args.batch_size392 hparams.sample_rate = args.sample_rate393 hparams.if_f0 = args.if_f0394 hparams.if_latest = args.if_latest395 hparams.save_every_weights = args.save_every_weights396 hparams.if_cache_data_in_gpu = args.if_cache_data_in_gpu397 hparams.data.training_files = "%s/filelist.txt" % experiment_dir398 399 hparams.train.log_interval = args.log_interval400 401 # Update log_interval in the 'train' section of the config dictionary402 config["train"]["log_interval"] = args.log_interval403 404 # Save the updated config back to the config_save_path405 with open(config_save_path, "w") as f:406 json.dump(config, f, indent=4)407 408 return hparams409 410 411def get_hparams_from_dir(model_dir):412 config_save_path = os.path.join(model_dir, "config.json")413 with open(config_save_path, "r") as f:414 data = f.read()415 config = json.loads(data)416 417 hparams = HParams(**config)418 hparams.model_dir = model_dir419 return hparams420 421 422def get_hparams_from_file(config_path):423 with open(config_path, "r") as f:424 data = f.read()425 config = json.loads(data)426 427 hparams = HParams(**config)428 return hparams429 430 431def check_git_hash(model_dir):432 source_dir = os.path.dirname(os.path.realpath(__file__))433 if not os.path.exists(os.path.join(source_dir, ".git")):434 logger.warn(435 "{} is not a git repository, therefore hash value comparison will be ignored.".format(436 source_dir437 )438 )439 return440 441 cur_hash = subprocess.getoutput("git rev-parse HEAD")442 443 path = os.path.join(model_dir, "githash")444 if os.path.exists(path):445 saved_hash = open(path).read()446 if saved_hash != cur_hash:447 logger.warn(448 "git hash values are different. {}(saved) != {}(current)".format(449 saved_hash[:8], cur_hash[:8]450 )451 )452 else:453 open(path, "w").write(cur_hash)454 455 456def get_logger(model_dir, filename="train.log"):457 global logger458 logger = logging.getLogger(os.path.basename(model_dir))459 logger.setLevel(logging.DEBUG)460 461 formatter = logging.Formatter("%(asctime)s\t%(name)s\t%(levelname)s\t%(message)s")462 if not os.path.exists(model_dir):463 os.makedirs(model_dir)464 h = logging.FileHandler(os.path.join(model_dir, filename))465 h.setLevel(logging.DEBUG)466 h.setFormatter(formatter)467 logger.addHandler(h)468 return logger469 470 471class HParams:472 def __init__(self, **kwargs):473 for k, v in kwargs.items():474 if type(v) == dict:475 v = HParams(**v)476 self[k] = v477 478 def keys(self):479 return self.__dict__.keys()480 481 def items(self):482 return self.__dict__.items()483 484 def values(self):485 return self.__dict__.values()486 487 def __len__(self):488 return len(self.__dict__)489 490 def __getitem__(self, key):491 return getattr(self, key)492 493 def __setitem__(self, key, value):494 return setattr(self, key, value)495 496 def __contains__(self, key):497 return key in self.__dict__498 499 def __repr__(self):500 return self.__dict__.__repr__()501