def __init__(self, n_tiers: int, layers: List[int], hidden_size: int, gmm_size: int, freq: int): """ Args: n_tiers (int): number of tiers the model is composed of layers (List[int]): list with the layers of every tier hidden_size (int): parameter for the hidden_state of the Delayed Stack Layers and other and other sizes gmm_size (int): number of mixture components of the GMM freq (int): size of the frequency axis of the spectrogram to generate. See note in the documentation of the file. """ super(MelNet, self).__init__() self.n_tiers = n_tiers self.layers = layers self.hidden_size = hidden_size self.gmm_size = gmm_size assert freq >= 2 ** (self.n_tiers / 2), "Size of frequency axis is too small for " \ "being generated with the number of tiers " \ "of this model" self.freq = freq self.tiers = nn.ModuleList([ Tier1( tier=1, n_layers=layers[0], hidden_size=hidden_size, gmm_size=gmm_size, # Calculate size of FREQ dimension for this tier freq=tierutil.get_size_freqdim_of_tier( n_mels=self.freq, n_tiers=self.n_tiers, tier=1)) ] + [ Tier( tier=tier_idx, n_layers=layers[tier_idx], hidden_size=hidden_size, gmm_size=gmm_size, # Calculate size of FREQ dimension for this tier freq=tierutil.get_size_freqdim_of_tier( n_mels=self.freq, n_tiers=self.n_tiers, tier=tier_idx + 1)) for tier_idx in range(1, n_tiers) ])
def sample(self, hp: HParams, synthesisp: HParams, timestamp: str, logger: logging.Logger, n_samples: int, length: int) -> torch.Tensor: """ Generates n_samples of audio of the given length. Args: hp (HParams): parameters. Parameters needed are hp.device synthesisp (HParams): parameters for performing the synthesis. Parameters needed are synthesisp.output_path to save the spectrogram generated at each tier. timestamp (str): information that identifies completely this run (synthesis). logger (logging.Logger): n_samples (int): amount of samples to generate. length (int): length of the samples to generate (in timesteps). Returns: spectrograms (torch.Tensor): samples of audio in spectrogram representation. Shape: [B=n_samples, FREQ=self.freq, FRAMES=length]. """ assert length >= 2 ** ( self.n_tiers / 2), "Length is too short for being generated with the " \ "number of tiers of this model." # Initially, the spectrogram (x) to generate it does not exist. x = None # --- TIER 1 ---- # The spectrogram is generated autoregressively, frame (length, or timestep) by frame. logger.info(f"Starting Tier 1/{self.n_tiers}") freq_of_tier1 = tierutil.get_size_freqdim_of_tier(n_mels=self.freq, n_tiers=self.n_tiers, tier=1) length_of_tier1 = tierutil.get_size_timedim_of_tier(timesteps=length, n_tiers=self.n_tiers, tier=1) for i in range(0, length_of_tier1): logger.info(f"Tier 1/{self.n_tiers} - Frame {i}/{length_of_tier1}") if x is None: # If the spectrogram has not been initialized, we initialized to an initial frame # of all zeros x = torch.zeros((n_samples, freq_of_tier1, 1), device=hp.device) else: # If the spectrogram has already been initialized, we have already computed some # frames. We concatenate a new frame initialized to all zeros which will be replaced # pixel by pixel by the new values # We change the shape from [B, FREQ, FRAMES] to [B, FREQ, FRAMES+1] by adding a new # frame x = torch.cat( [x, torch.zeros((n_samples, freq_of_tier1, 1), device=hp.device)], dim=-1) # Inside a frame, the spectrogram is generated autoregressively, freq by freq for j in range(0, freq_of_tier1): # we generate the parameters for all the spectrogram (across all samples) mu_hat, std_hat, pi_hat = self.tiers[0](x) # with the parameters we generate the values of the next spectrogram # (across all samples) new_spectrogram = sample_gmm_batch(mu_hat, std_hat, pi_hat) # but only use the value of the new pixel that we are generating # (across all samples) since the spectrogram is generated autoregressively x[:, j, i] = new_spectrogram[:, j, i] # Save spectrogram generated at tier1 torch.save(x, f"{synthesisp.output_path}/{timestamp}_tier1.pt") # --- TIER >1 --- for tier_idx in range(2, self.n_tiers + 1): temp_x = None # temporary spectrogram that will be generated by this tier # The spectrogram is generated autoregressively, frame (length, or timestep) by frame. logger.info(f"Starting Tier {tier_idx}/{self.n_tiers}") freq_of_tierX = tierutil.get_size_freqdim_of_tier(n_mels=self.freq, n_tiers=self.n_tiers, tier=tier_idx) length_of_tierX = tierutil.get_size_timedim_of_tier(timesteps=length, n_tiers=self.n_tiers, tier=tier_idx) for i in range(0, length_of_tierX): logger.info(f"Tier {tier_idx}/{self.n_tiers} - Frame {i}/{length_of_tierX}") if temp_x is None: # If the spectrogram of this tier has not been initialized, we initialized to an # initial frame of all zeros temp_x = torch.zeros((n_samples, freq_of_tierX, 1), device=hp.device) else: # If the spectrogram of this tier has already been initialized, we have already # computed some frames. We concatenate a new frame initialized to all zeros # which will be replaced pixel by pixel by the new values # We change shape from [B, FREQ, FRAMES] to [B, FREQ, FRAMES+1] by adding a new # frame temp_x = torch.cat( [temp_x, torch.zeros((n_samples, freq_of_tierX, 1), device=hp.device)], dim=-1) # Inside a frame, the spectrogram is generated autoregressively, freq by freq for j in range(0, freq_of_tierX): # we generate the parameters for all the spectrogram (across all samples) mu_hat, std_hat, pi_hat = self.tiers[tier_idx - 1](temp_x, x) # with the parameters we generate the values of the next spectrogram # (across all samples) new_spectrogram = sample_gmm_batch(mu_hat, std_hat, pi_hat) # but only use the value of the new pixel that we are generating # (across all samples) since the spectrogram is generated autoregressively temp_x[:, j, i] = new_spectrogram[:, j, i] # After generating the spectrogram of this tier, we interleave it to put it together # with the spectrogram generated by previous tiers. In the next iteration, this will # be the input to condition the next tier x = tierutil.interleave(temp_x, x, tier_idx) x = x.to(hp.device) # Save spectrogram generated at tier1 torch.save(temp_x, f"{synthesisp.output_path}/{timestamp}_tier{tier_idx}.pt") torch.save(x, f"{synthesisp.output_path}/{timestamp}_tier1-tier{tier_idx}.pt") return x
def train_tier(args: argparse.Namespace, hp: HParams, tier: int, extension_architecture: str, timestamp: str, tensorboardwriter: TensorboardWriter, logger: logging.Logger) -> None: """ Trains one tier of MelNet. Args: args (argparse.Namespace): parameters to set up the training. At least, args must contain: args = {"path_config": ..., "tier": ..., "checkpoint_path": ...} hp (HParams): hyperparameters for the model and other parameters (training, dataset, ...) tier (int): number of the tier to train. extension_architecture (str): information about the network's architecture of this run (training) to identify the logs and weights of the model. timestamp (str): information that identifies completely this run (training). tensorboardwriter (TensorboardWriter): to log information about training to tensorboard. logger (logging.Logger): to log general information about the training of the model. """ logger.info(f"Start training of tier {tier}/{hp.network.n_tiers}") # Setup the data ready to be consumed train_dataloader, test_dataloader, num_samples = get_dataloader(hp) # Setup tier # Calculate size of FREQ dimension for this tier tier_freq = tierutil.get_size_freqdim_of_tier(n_mels=hp.audio.mel_channels, n_tiers=hp.network.n_tiers, tier=tier) if tier == 1: model = Tier1(tier=tier, n_layers=hp.network.layers[tier - 1], hidden_size=hp.network.hidden_size, gmm_size=hp.network.gmm_size, freq=tier_freq) else: model = Tier(tier=tier, n_layers=hp.network.layers[tier - 1], hidden_size=hp.network.hidden_size, gmm_size=hp.network.gmm_size, freq=tier_freq) model = model.to(hp.device) model.train() parameters = model.parameters() # Setup loss criterion and optimizer criterion = GMMLoss() optimizer = torch.optim.RMSprop(params=parameters, lr=hp.training.lr, momentum=hp.training.momentum) # Check if training has to be resumed from previous checkpoint if args.checkpoint_path is not None: model, optimizer = resume_training(args, hp, tier, model, optimizer, logger) else: logger.info( f"Starting new training on dataset {hp.data.dataset} with configuration file " f"name {hp.name}") # Train the tier total_iterations = 0 loss_logging = 0 # accumulated loss between logging iterations loss_save = 0 # accumulated loss between saving iterations prev_loss_onesample = 1e8 # used to compare between saving iterations and decide whether or not # to save the model gradients = [] for epoch in range(hp.training.epochs): logger.info(f"Epoch: {epoch}/{hp.training.epochs} - Starting") for i, (waveform, utterance) in enumerate(train_dataloader): # 1.1 Transform waveform input to melspectrogram and apply preprocessing to normalize waveform = waveform.to(device=hp.device, non_blocking=True) spectrogram = transforms.wave_to_melspectrogram(waveform, hp) spectrogram = audio_normalizing.preprocessing(spectrogram, hp) # 1.2 Get input and output from the original spectrogram for this tier input_spectrogram, output_spectrogram = tierutil.split( spectrogram=spectrogram, tier=tier, n_tiers=hp.network.n_tiers) length_spectrogram = input_spectrogram.size(2) # if item is too long, we jump to the next one if length_spectrogram > 1000: continue # 2. Compute the model output if tier == 1: # generation is unconditional so there is only one input mu_hat, std_hat, pi_hat = model(spectrogram=input_spectrogram) else: # generation is conditional on the spectrogram generated by previous tiers mu_hat, std_hat, pi_hat = model( spectrogram=output_spectrogram, spectrogram_prev_tier=input_spectrogram) # gpumemory.stat_cuda("Forward") # 3. Calculate the loss loss = criterion(mu=mu_hat, std=std_hat, pi=pi_hat, target=output_spectrogram) # gpumemory.stat_cuda("Loss") del spectrogram del mu_hat, std_hat, pi_hat # 3.1 Check if loss has exploded if torch.isnan(loss) or torch.isinf(loss): error_msg = f"Loss exploded at Epoch: {epoch}/{hp.training.epochs} - " \ f"Iteration: {i * hp.training.batch_size}/{num_samples}" logger.error(error_msg) raise Exception(error_msg) # 4. Compute gradients loss_cpu = loss.item() loss = loss / hp.training.accumulation_steps loss.backward() # 5. Perform backpropagation (using gradient accumulation so efective batch size is the # same as in the paper) if (total_iterations + 1) % (hp.training.accumulation_steps / hp.training.batch_size) == 0: gradients.append(gradient_norm(model)) avg_gradient = sum(gradients) / len(gradients) logger.info(f"Gradient norm: {gradients[-1]} - " f"Avg gradient: {avg_gradient}") torch.nn.utils.clip_grad_norm_(parameters, 2200) optimizer.step() model.zero_grad() # 6. Logging and saving model loss_oneframe = loss_cpu / (length_spectrogram * hp.training.batch_size) loss_logging += loss_oneframe # accumulated loss between logging iterations loss_save += loss_oneframe # accumulated loss between saving iterations # 6.1 Save model (if is better than previous tier) if (total_iterations + 1) % hp.training.save_iterations == 0: # Calculate average loss of one sample of a batch loss_onesample = loss_save / hp.training.save_iterations # if loss_onesample of these iterations is lower, the tier is better and we save it if loss_onesample <= prev_loss_onesample: path = f"{hp.training.dir_chkpt}/tier{tier}_{timestamp}_loss{loss_onesample:.2f}.pt" torch.save(obj={ 'dataset': hp.data.dataset, 'tier_idx': tier, 'hp': hp, 'epoch': epoch, 'iterations': i, 'total_iterations': total_iterations, 'tier': model.state_dict(), 'optimizer': optimizer.state_dict() }, f=path) logger.info(f"Model saved to: {path}") prev_loss_onesample = loss_onesample loss_save = 0 # 6.2 Logging if (total_iterations + 1) % hp.logging.log_iterations == 0: # Calculate average loss of one sample of a batch loss_onesample = loss_logging / hp.logging.log_iterations tensorboardwriter.log_training(hp, loss_onesample, total_iterations) logger.info( f"Epoch: {epoch}/{hp.training.epochs} - " f"Iteration: {i * hp.training.batch_size}/{num_samples} - " f"Loss: {loss_onesample:.4f}") loss_logging = 0 # 6.3 Evaluate if (total_iterations + 1) % hp.training.evaluation_iterations == 0: evaluation(hp, tier, test_dataloader, model, criterion, logger) total_iterations += 1 # After finishing training: save model, hyperparameters and total loss path = f"{hp.training.dir_chkpt}/tier{tier}_{timestamp}_epoch{epoch}_final.pt" torch.save(obj={ 'dataset': hp.data.dataset, 'tier_idx': tier, 'hp': hp, 'epoch': epoch, 'iterations': evaluation(hp, tier, test_dataloader, model, criterion, logger), 'total_iterations': total_iterations, 'tier': model.state_dict(), 'optimizer': optimizer.state_dict() }, f=path) logger.info(f"Model saved to: {path}") tensorboardwriter.log_end_training(hp=hp, loss=-1) logger.info("Finished training")
def sample(self, hp: HParams, synthesisp: HParams, timestamp: str, logger: logging.Logger, n_samples: int, length: int) -> torch.Tensor: """ Generates n_samples of audio of the given length. Args: hp (HParams): parameters. Parameters needed are hp.device synthesisp (HParams): parameters for performing the synthesis. Parameters needed are synthesisp.output_path to save the spectrogram generated at each tier. timestamp (str): information that identifies completely this run (synthesis). logger (logging.Logger): n_samples (int): amount of samples to generate. length (int): length of the samples to generate (in timesteps). Returns: spectrograms (torch.Tensor): samples of audio in spectrogram representation. Shape: [B=n_samples, FREQ=self.freq, FRAMES=length]. """ assert length >= 2 ** ( self.n_tiers / 2), "Length is too short for being generated with the " \ "number of tiers of this model." # Initially, the spectrogram (x) to generate it does not exist. x = None # Load a spectrogram from the dataset from src.utils.training_batch import get_dataloader from src.dataprocessing import transforms as T from src.dataprocessing.audio_normalizing import preprocessing dataloader, _, _ = get_dataloader(hp) wave = None for i, (waveform, utterance) in enumerate(dataloader): if "building" in utterance[0]: print(utterance[0]) if utterance[ 0] == "One building, Market Hall, was unavailable for November 22.": wave = waveform break if wave is None: logger.info("wave not found") return #dataiter = iter(dataloader) #wave, utterance = dataiter.next() waveform = wave.to(device=hp.device, non_blocking=True) spectrogram = T.wave_to_melspectrogram(waveform, hp) spectrogram = preprocessing(spectrogram, hp) # Split the spectrogram to get the spectrogram that would be the output of the first tier input_spectrogram, output_spectrogram = tierutil.split( spectrogram=spectrogram, tier=1, n_tiers=hp.network.n_tiers) # Use the spectrogram from the dataset as output from first tier x = output_spectrogram length = spectrogram.size(2) # Save spectrogram generated at tier1 torch.save(x, f"{synthesisp.output_path}/{timestamp}_tier1.pt") # --- TIER >1 --- for tier_idx in range(2, self.n_tiers + 1): temp_x = None # temporary spectrogram that will be generated by this tier # The spectrogram is generated autoregressively, frame (length, or timestep) by frame. logger.info(f"Starting Tier {tier_idx}/{self.n_tiers}") freq_of_tierX = tierutil.get_size_freqdim_of_tier( n_mels=self.freq, n_tiers=self.n_tiers, tier=tier_idx) length_of_tierX = tierutil.get_size_timedim_of_tier( timesteps=length, n_tiers=self.n_tiers, tier=tier_idx) print("Shape of original spectrogram: ", spectrogram.size()) print("Shape of spectrogram_prev_tier (x): ", x.size()) print("Freq_of_tierX: ", freq_of_tierX) print("Length_of_tierX: ", length_of_tierX) length_of_tierX = min(length_of_tierX, x.size(2)) x = x[:, :, :length_of_tierX] for i in range(0, length_of_tierX): logger.info( f"Tier {tier_idx}/{self.n_tiers} - Frame {i}/{length_of_tierX}" ) if temp_x is None: # If the spectrogram of this tier has not been initialized, we initialized to an # initial frame of all zeros temp_x = torch.zeros((n_samples, freq_of_tierX, 1), device=hp.device) else: # If the spectrogram of this tier has already been initialized, we have already # computed some frames. We concatenate a new frame initialized to all zeros # which will be replaced pixel by pixel by the new values # We change shape from [B, FREQ, FRAMES] to [B, FREQ, FRAMES+1] by adding a new # frame temp_x = torch.cat([ temp_x, torch.zeros( (n_samples, freq_of_tierX, 1), device=hp.device) ], dim=-1) # Inside a frame, the spectrogram is generated autoregressively, freq by freq for j in range(0, freq_of_tierX): # we generate the parameters for all the spectrogram (across all samples) mu_hat, std_hat, pi_hat = self.tiers[tier_idx - 1](temp_x, x) # with the parameters we generate the values of the next spectrogram # (across all samples) new_spectrogram = sample_gmm_batch(mu_hat, std_hat, pi_hat) # but only use the value of the new pixel that we are generating # (across all samples) since the spectrogram is generated autoregressively temp_x[:, j, i] = new_spectrogram[:, j, i] # After generating the spectrogram of this tier, we interleave it to put it together # with the spectrogram generated by previous tiers. In the next iteration, this will # be the input to condition the next tier x = tierutil.interleave(temp_x, x, tier_idx) x = x.to(hp.device) # Save spectrogram generated at tier1 torch.save( temp_x, f"{synthesisp.output_path}/{timestamp}_tier{tier_idx}.pt") torch.save( x, f"{synthesisp.output_path}/{timestamp}_tier1-tier{tier_idx}.pt" ) return x