Spaces:
Runtime error
Runtime error
| import os, traceback | |
| import glob | |
| import sys | |
| import argparse | |
| import logging | |
| import json | |
| import subprocess | |
| import numpy as np | |
| from scipy.io.wavfile import read | |
| import torch | |
| MATPLOTLIB_FLAG = False | |
| logging.basicConfig(stream=sys.stdout, level=logging.DEBUG) | |
| logger = logging | |
| def load_checkpoint_d(checkpoint_path, combd, sbd, optimizer=None, load_opt=1): | |
| assert os.path.isfile(checkpoint_path) | |
| checkpoint_dict = torch.load(checkpoint_path, map_location="cpu") | |
| ################## | |
| def go(model, bkey): | |
| saved_state_dict = checkpoint_dict[bkey] | |
| if hasattr(model, "module"): | |
| state_dict = model.module.state_dict() | |
| else: | |
| state_dict = model.state_dict() | |
| new_state_dict = {} | |
| for k, v in state_dict.items(): # 模型需要的shape | |
| try: | |
| new_state_dict[k] = saved_state_dict[k] | |
| if saved_state_dict[k].shape != state_dict[k].shape: | |
| print( | |
| "shape-%s-mismatch|need-%s|get-%s" | |
| % (k, state_dict[k].shape, saved_state_dict[k].shape) | |
| ) # | |
| raise KeyError | |
| except: | |
| # logger.info(traceback.format_exc()) | |
| logger.info("%s is not in the checkpoint" % k) # pretrain缺失的 | |
| new_state_dict[k] = v # 模型自带的随机值 | |
| if hasattr(model, "module"): | |
| model.module.load_state_dict(new_state_dict, strict=False) | |
| else: | |
| model.load_state_dict(new_state_dict, strict=False) | |
| go(combd, "combd") | |
| go(sbd, "sbd") | |
| ############# | |
| logger.info("Loaded model weights") | |
| iteration = checkpoint_dict["iteration"] | |
| learning_rate = checkpoint_dict["learning_rate"] | |
| if ( | |
| optimizer is not None and load_opt == 1 | |
| ): ###加载不了，如果是空的的话，重新初始化，可能还会影响lr时间表的更新，因此在train文件最外围catch | |
| # try: | |
| optimizer.load_state_dict(checkpoint_dict["optimizer"]) | |
| # except: | |
| # traceback.print_exc() | |
| logger.info("Loaded checkpoint '{}' (epoch {})".format(checkpoint_path, iteration)) | |
| return model, optimizer, learning_rate, iteration | |
| # def load_checkpoint(checkpoint_path, model, optimizer=None): | |
| # assert os.path.isfile(checkpoint_path) | |
| # checkpoint_dict = torch.load(checkpoint_path, map_location='cpu') | |
| # iteration = checkpoint_dict['iteration'] | |
| # learning_rate = checkpoint_dict['learning_rate'] | |
| # if optimizer is not None: | |
| # optimizer.load_state_dict(checkpoint_dict['optimizer']) | |
| # # print(1111) | |
| # saved_state_dict = checkpoint_dict['model'] | |
| # # print(1111) | |
| # | |
| # if hasattr(model, 'module'): | |
| # state_dict = model.module.state_dict() | |
| # else: | |
| # state_dict = model.state_dict() | |
| # new_state_dict= {} | |
| # for k, v in state_dict.items(): | |
| # try: | |
| # new_state_dict[k] = saved_state_dict[k] | |
| # except: | |
| # logger.info("%s is not in the checkpoint" % k) | |
| # new_state_dict[k] = v | |
| # if hasattr(model, 'module'): | |
| # model.module.load_state_dict(new_state_dict) | |
| # else: | |
| # model.load_state_dict(new_state_dict) | |
| # logger.info("Loaded checkpoint '{}' (epoch {})" .format( | |
| # checkpoint_path, iteration)) | |
| # return model, optimizer, learning_rate, iteration | |
| def load_checkpoint(checkpoint_path, model, optimizer=None, load_opt=1): | |
| assert os.path.isfile(checkpoint_path) | |
| checkpoint_dict = torch.load(checkpoint_path, map_location="cpu") | |
| saved_state_dict = checkpoint_dict["model"] | |
| if hasattr(model, "module"): | |
| state_dict = model.module.state_dict() | |
| else: | |
| state_dict = model.state_dict() | |
| new_state_dict = {} | |
| for k, v in state_dict.items(): # 模型需要的shape | |
| try: | |
| new_state_dict[k] = saved_state_dict[k] | |
| if saved_state_dict[k].shape != state_dict[k].shape: | |
| print( | |
| "shape-%s-mismatch|need-%s|get-%s" | |
| % (k, state_dict[k].shape, saved_state_dict[k].shape) | |
| ) # | |
| raise KeyError | |
| except: | |
| # logger.info(traceback.format_exc()) | |
| logger.info("%s is not in the checkpoint" % k) # pretrain缺失的 | |
| new_state_dict[k] = v # 模型自带的随机值 | |
| if hasattr(model, "module"): | |
| model.module.load_state_dict(new_state_dict, strict=False) | |
| else: | |
| model.load_state_dict(new_state_dict, strict=False) | |
| logger.info("Loaded model weights") | |
| iteration = checkpoint_dict["iteration"] | |
| learning_rate = checkpoint_dict["learning_rate"] | |
| if ( | |
| optimizer is not None and load_opt == 1 | |
| ): ###加载不了，如果是空的的话，重新初始化，可能还会影响lr时间表的更新，因此在train文件最外围catch | |
| # try: | |
| optimizer.load_state_dict(checkpoint_dict["optimizer"]) | |
| # except: | |
| # traceback.print_exc() | |
| logger.info("Loaded checkpoint '{}' (epoch {})".format(checkpoint_path, iteration)) | |
| return model, optimizer, learning_rate, iteration | |
| def save_checkpoint(model, optimizer, learning_rate, iteration, checkpoint_path): | |
| logger.info( | |
| "Saving model and optimizer state at epoch {} to {}".format( | |
| iteration, checkpoint_path | |
| ) | |
| ) | |
| if hasattr(model, "module"): | |
| state_dict = model.module.state_dict() | |
| else: | |
| state_dict = model.state_dict() | |
| torch.save( | |
| { | |
| "model": state_dict, | |
| "iteration": iteration, | |
| "optimizer": optimizer.state_dict(), | |
| "learning_rate": learning_rate, | |
| }, | |
| checkpoint_path, | |
| ) | |
| def save_checkpoint_d(combd, sbd, optimizer, learning_rate, iteration, checkpoint_path): | |
| logger.info( | |
| "Saving model and optimizer state at epoch {} to {}".format( | |
| iteration, checkpoint_path | |
| ) | |
| ) | |
| if hasattr(combd, "module"): | |
| state_dict_combd = combd.module.state_dict() | |
| else: | |
| state_dict_combd = combd.state_dict() | |
| if hasattr(sbd, "module"): | |
| state_dict_sbd = sbd.module.state_dict() | |
| else: | |
| state_dict_sbd = sbd.state_dict() | |
| torch.save( | |
| { | |
| "combd": state_dict_combd, | |
| "sbd": state_dict_sbd, | |
| "iteration": iteration, | |
| "optimizer": optimizer.state_dict(), | |
| "learning_rate": learning_rate, | |
| }, | |
| checkpoint_path, | |
| ) | |
| def summarize( | |
| writer, | |
| global_step, | |
| scalars={}, | |
| histograms={}, | |
| images={}, | |
| audios={}, | |
| audio_sampling_rate=22050, | |
| ): | |
| for k, v in scalars.items(): | |
| writer.add_scalar(k, v, global_step) | |
| for k, v in histograms.items(): | |
| writer.add_histogram(k, v, global_step) | |
| for k, v in images.items(): | |
| writer.add_image(k, v, global_step, dataformats="HWC") | |
| for k, v in audios.items(): | |
| writer.add_audio(k, v, global_step, audio_sampling_rate) | |
| def latest_checkpoint_path(dir_path, regex="G_*.pth"): | |
| f_list = glob.glob(os.path.join(dir_path, regex)) | |
| f_list.sort(key=lambda f: int("".join(filter(str.isdigit, f)))) | |
| x = f_list[-1] | |
| print(x) | |
| return x | |
| def plot_spectrogram_to_numpy(spectrogram): | |
| global MATPLOTLIB_FLAG | |
| if not MATPLOTLIB_FLAG: | |
| import matplotlib | |
| matplotlib.use("Agg") | |
| MATPLOTLIB_FLAG = True | |
| mpl_logger = logging.getLogger("matplotlib") | |
| mpl_logger.setLevel(logging.WARNING) | |
| import matplotlib.pylab as plt | |
| import numpy as np | |
| fig, ax = plt.subplots(figsize=(10, 2)) | |
| im = ax.imshow(spectrogram, aspect="auto", origin="lower", interpolation="none") | |
| plt.colorbar(im, ax=ax) | |
| plt.xlabel("Frames") | |
| plt.ylabel("Channels") | |
| plt.tight_layout() | |
| fig.canvas.draw() | |
| data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep="") | |
| data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,)) | |
| plt.close() | |
| return data | |
| def plot_alignment_to_numpy(alignment, info=None): | |
| global MATPLOTLIB_FLAG | |
| if not MATPLOTLIB_FLAG: | |
| import matplotlib | |
| matplotlib.use("Agg") | |
| MATPLOTLIB_FLAG = True | |
| mpl_logger = logging.getLogger("matplotlib") | |
| mpl_logger.setLevel(logging.WARNING) | |
| import matplotlib.pylab as plt | |
| import numpy as np | |
| fig, ax = plt.subplots(figsize=(6, 4)) | |
| im = ax.imshow( | |
| alignment.transpose(), aspect="auto", origin="lower", interpolation="none" | |
| ) | |
| fig.colorbar(im, ax=ax) | |
| xlabel = "Decoder timestep" | |
| if info is not None: | |
| xlabel += "\n\n" + info | |
| plt.xlabel(xlabel) | |
| plt.ylabel("Encoder timestep") | |
| plt.tight_layout() | |
| fig.canvas.draw() | |
| data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep="") | |
| data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,)) | |
| plt.close() | |
| return data | |
| def load_wav_to_torch(full_path): | |
| sampling_rate, data = read(full_path) | |
| return torch.FloatTensor(data.astype(np.float32)), sampling_rate | |
| def load_filepaths_and_text(filename, split="|"): | |
| with open(filename, encoding='utf-8') as f: | |
| filepaths_and_text = [line.strip().split(split) for line in f] | |
| filepaths_and_text = [item for item in filepaths_and_text if len(item) == 5] # ensure there are 5 items. | |
| return filepaths_and_text | |
| def get_hparams(init=True): | |
| """ | |
| todo: | |
| 结尾七人组： | |
| 保存频率、总epoch done | |
| bs done | |
| pretrainG、pretrainD done | |
| 卡号：os.en["CUDA_VISIBLE_DEVICES"] done | |
| if_latest done | |
| 模型：if_f0 done | |
| 采样率：自动选择config done | |
| 是否缓存数据集进GPU:if_cache_data_in_gpu done | |
| -m: | |
| 自动决定training_files路径,改掉train_nsf_load_pretrain.py里的hps.data.training_files done | |
| -c不要了 | |
| """ | |
| parser = argparse.ArgumentParser() | |
| # parser.add_argument('-c', '--config', type=str, default="configs/40k.json",help='JSON file for configuration') | |
| parser.add_argument( | |
| "-se", | |
| "--save_every_epoch", | |
| type=int, | |
| required=True, | |
| help="checkpoint save frequency (epoch)", | |
| ) | |
| parser.add_argument( | |
| "-te", "--total_epoch", type=int, required=True, help="total_epoch" | |
| ) | |
| parser.add_argument( | |
| "-pg", "--pretrainG", type=str, default="", help="Pretrained Discriminator path" | |
| ) | |
| parser.add_argument( | |
| "-pd", "--pretrainD", type=str, default="", help="Pretrained Generator path" | |
| ) | |
| parser.add_argument("-g", "--gpus", type=str, default="0", help="split by -") | |
| parser.add_argument( | |
| "-bs", "--batch_size", type=int, required=True, help="batch size" | |
| ) | |
| parser.add_argument( | |
| "-e", "--experiment_dir", type=str, required=True, help="experiment dir" | |
| ) # -m | |
| parser.add_argument( | |
| "-sr", "--sample_rate", type=str, required=True, help="sample rate, 32k/40k/48k" | |
| ) | |
| parser.add_argument( | |
| "-sw", | |
| "--save_every_weights", | |
| type=str, | |
| default="0", | |
| help="save the extracted model in weights directory when saving checkpoints", | |
| ) | |
| parser.add_argument( | |
| "-v", "--version", type=str, required=True, help="model version" | |
| ) | |
| parser.add_argument( | |
| "-f0", | |
| "--if_f0", | |
| type=int, | |
| required=True, | |
| help="use f0 as one of the inputs of the model, 1 or 0", | |
| ) | |
| parser.add_argument( | |
| "-l", | |
| "--if_latest", | |
| type=int, | |
| required=True, | |
| help="if only save the latest G/D pth file, 1 or 0", | |
| ) | |
| parser.add_argument( | |
| "-c", | |
| "--if_cache_data_in_gpu", | |
| type=int, | |
| required=True, | |
| help="if caching the dataset in GPU memory, 1 or 0", | |
| ) | |
| parser.add_argument( | |
| "-li", "--log_interval", type=int, required=True, help="log interval" | |
| ) | |
| args = parser.parse_args() | |
| name = args.experiment_dir | |
| experiment_dir = os.path.join("./logs", args.experiment_dir) | |
| if not os.path.exists(experiment_dir): | |
| os.makedirs(experiment_dir) | |
| if args.version == "v1" or args.sample_rate == "40k": | |
| config_path = "configs/%s.json" % args.sample_rate | |
| else: | |
| config_path = "configs/%s_v2.json" % args.sample_rate | |
| config_save_path = os.path.join(experiment_dir, "config.json") | |
| if init: | |
| with open(config_path, "r") as f: | |
| data = f.read() | |
| with open(config_save_path, "w") as f: | |
| f.write(data) | |
| else: | |
| with open(config_save_path, "r") as f: | |
| data = f.read() | |
| config = json.loads(data) | |
| hparams = HParams(**config) | |
| hparams.model_dir = hparams.experiment_dir = experiment_dir | |
| hparams.save_every_epoch = args.save_every_epoch | |
| hparams.name = name | |
| hparams.total_epoch = args.total_epoch | |
| hparams.pretrainG = args.pretrainG | |
| hparams.pretrainD = args.pretrainD | |
| hparams.version = args.version | |
| hparams.gpus = args.gpus | |
| hparams.train.batch_size = args.batch_size | |
| hparams.sample_rate = args.sample_rate | |
| hparams.if_f0 = args.if_f0 | |
| hparams.if_latest = args.if_latest | |
| hparams.save_every_weights = args.save_every_weights | |
| hparams.if_cache_data_in_gpu = args.if_cache_data_in_gpu | |
| hparams.data.training_files = "%s/filelist.txt" % experiment_dir | |
| hparams.train.log_interval = args.log_interval | |
| # Update log_interval in the 'train' section of the config dictionary | |
| config["train"]["log_interval"] = args.log_interval | |
| # Save the updated config back to the config_save_path | |
| with open(config_save_path, "w") as f: | |
| json.dump(config, f, indent=4) | |
| return hparams | |
| def get_hparams_from_dir(model_dir): | |
| config_save_path = os.path.join(model_dir, "config.json") | |
| with open(config_save_path, "r") as f: | |
| data = f.read() | |
| config = json.loads(data) | |
| hparams = HParams(**config) | |
| hparams.model_dir = model_dir | |
| return hparams | |
| def get_hparams_from_file(config_path): | |
| with open(config_path, "r") as f: | |
| data = f.read() | |
| config = json.loads(data) | |
| hparams = HParams(**config) | |
| return hparams | |
| def check_git_hash(model_dir): | |
| source_dir = os.path.dirname(os.path.realpath(__file__)) | |
| if not os.path.exists(os.path.join(source_dir, ".git")): | |
| logger.warn( | |
| "{} is not a git repository, therefore hash value comparison will be ignored.".format( | |
| source_dir | |
| ) | |
| ) | |
| return | |
| cur_hash = subprocess.getoutput("git rev-parse HEAD") | |
| path = os.path.join(model_dir, "githash") | |
| if os.path.exists(path): | |
| saved_hash = open(path).read() | |
| if saved_hash != cur_hash: | |
| logger.warn( | |
| "git hash values are different. {}(saved) != {}(current)".format( | |
| saved_hash[:8], cur_hash[:8] | |
| ) | |
| ) | |
| else: | |
| open(path, "w").write(cur_hash) | |
| def get_logger(model_dir, filename="train.log"): | |
| global logger | |
| logger = logging.getLogger(os.path.basename(model_dir)) | |
| logger.setLevel(logging.DEBUG) | |
| formatter = logging.Formatter("%(asctime)s\t%(name)s\t%(levelname)s\t%(message)s") | |
| if not os.path.exists(model_dir): | |
| os.makedirs(model_dir) | |
| h = logging.FileHandler(os.path.join(model_dir, filename)) | |
| h.setLevel(logging.DEBUG) | |
| h.setFormatter(formatter) | |
| logger.addHandler(h) | |
| return logger | |
| class HParams: | |
| def __init__(self, **kwargs): | |
| for k, v in kwargs.items(): | |
| if type(v) == dict: | |
| v = HParams(**v) | |
| self[k] = v | |
| def keys(self): | |
| return self.__dict__.keys() | |
| def items(self): | |
| return self.__dict__.items() | |
| def values(self): | |
| return self.__dict__.values() | |
| def __len__(self): | |
| return len(self.__dict__) | |
| def __getitem__(self, key): | |
| return getattr(self, key) | |
| def __setitem__(self, key, value): | |
| return setattr(self, key, value) | |
| def __contains__(self, key): | |
| return key in self.__dict__ | |
| def __repr__(self): | |
| return self.__dict__.__repr__() | |