Merge pull request #109 from TheButlah/master

Enabled multi-gpu training, buffers, grad clip in vocoder, saving optimizer state, and more fixes

Author: fatchord

.gitignore

@@ -1,5 +1,6 @@
-# PyCharm files
+# IDE files
 .idea
+.vscode
 
 # Mac files
 .DS_Store

gen_tacotron.py

@@ -91,7 +91,7 @@
 
     simple_table([('WaveRNN', str(voc_k) + 'k'),
                   ('Tacotron', str(tts_k) + 'k'),
-                  ('r', tts_model.r.item()),
+                  ('r', tts_model.r),
                   ('Generation Mode', 'Batched' if batched else 'Unbatched'),
                   ('Target Samples', target if batched else 'N/A'),
                   ('Overlap Samples', overlap if batched else 'N/A')])

gen_wavernn.py

@@ -7,7 +7,7 @@
 import argparse
 
 
-def gen_testset(model, test_set, samples, batched, target, overlap, save_path):
+def gen_testset(model: WaveRNN, test_set, samples, batched, target, overlap, save_path):
 
     k = model.get_step() // 1000
 
@@ -34,7 +34,7 @@ def gen_testset(model, test_set, samples, batched, target, overlap, save_path):
         _ = model.generate(m, save_str, batched, target, overlap, hp.mu_law)
 
 
-def gen_from_file(model, load_path, save_path, batched, target, overlap):
+def gen_from_file(model: WaveRNN, load_path, save_path, batched, target, overlap):
 
     k = model.get_step() // 1000
     file_name = load_path.split('/')[-1]
@@ -61,7 +61,7 @@ def gen_from_file(model, load_path, save_path, batched, target, overlap):
     parser.add_argument('--overlap', '-o', type=int, help='[int] number of crossover samples')
     parser.add_argument('--file', '-f', type=str, help='[string/path] for testing a wav outside dataset')
     parser.add_argument('--weights', '-w', type=str, help='[string/path] checkpoint file to load weights from')
-    parser.add_argument('--gta', '-g', dest='use_gta', action='store_true', help='Generate from GTA testset')
+    parser.add_argument('--gta', '-g', dest='gta', action='store_true', help='Generate from GTA testset')
     parser.add_argument('--force_cpu', '-c', action='store_true', help='Forces CPU-only training, even when in CUDA capable environment')
 
     parser.set_defaults(batched=hp.voc_gen_batched)

hparams.py

@@ -52,6 +52,7 @@
 voc_test_samples = 50               # How many unseen samples to put aside for testing
 voc_pad = 2                         # this will pad the input so that the resnet can 'see' wider than input length
 voc_seq_len = hop_length * 5        # must be a multiple of hop_length
+voc_clip_grad_norm = 4              # set to None if no gradient clipping needed
 
 # Generating / Synthesizing
 voc_gen_batched = True              # very fast (realtime+) single utterance batched generation

models/deepmind_version.py

@@ -3,6 +3,7 @@
 import torch.nn.functional as F
 from utils.display import *
 from utils.dsp import *
+import numpy as np
 
 class WaveRNN(nn.Module):
     def __init__(self, hidden_size=896, quantisation=256):
@@ -167,7 +168,9 @@ def get_initial_hidden(self, batch_size=1):
         device = next(self.parameters()).device  # use same device as parameters
         return torch.zeros(batch_size, self.hidden_size, device=device)
 
-    def num_params(self):
+    def num_params(self, print_out=True):
         parameters = filter(lambda p: p.requires_grad, self.parameters())
         parameters = sum([np.prod(p.size()) for p in parameters]) / 1_000_000
-        print('Trainable Parameters: %.3f million' % parameters)
+        if print_out:
+            print('Trainable Parameters: %.3f million' % parameters)
+        return parameters

models/fatchord_version.py

@@ -5,6 +5,7 @@
 from utils.display import *
 from utils.dsp import *
 import os
+import numpy as np
 
 
 class ResBlock(nn.Module):
@@ -100,24 +101,38 @@ def __init__(self, rnn_dims, fc_dims, bits, pad, upsample_factors,
         else:
             RuntimeError("Unknown model mode value - ", self.mode)
 
+        # List of rnns to call `flatten_parameters()` on
+        self._to_flatten = []
+        
         self.rnn_dims = rnn_dims
         self.aux_dims = res_out_dims // 4
         self.hop_length = hop_length
         self.sample_rate = sample_rate
 
         self.upsample = UpsampleNetwork(feat_dims, upsample_factors, compute_dims, res_blocks, res_out_dims, pad)
         self.I = nn.Linear(feat_dims + self.aux_dims + 1, rnn_dims)
+        
         self.rnn1 = nn.GRU(rnn_dims, rnn_dims, batch_first=True)
         self.rnn2 = nn.GRU(rnn_dims + self.aux_dims, rnn_dims, batch_first=True)
+        self._to_flatten += [self.rnn1, self.rnn2]
+        
         self.fc1 = nn.Linear(rnn_dims + self.aux_dims, fc_dims)
         self.fc2 = nn.Linear(fc_dims + self.aux_dims, fc_dims)
         self.fc3 = nn.Linear(fc_dims, self.n_classes)
 
-        self.step = nn.Parameter(torch.zeros(1).long(), requires_grad=False)
+        self.register_buffer('step', torch.zeros(1, dtype=torch.long))
         self.num_params()
 
+        # Avoid fragmentation of RNN parameters and associated warning
+        self._flatten_parameters()
+
     def forward(self, x, mels):
         device = next(self.parameters()).device  # use same device as parameters
+
+        # Although we `_flatten_parameters()` on init, when using DataParallel 
+        # the model gets replicated, making it no longer guaranteed that the
+        # weights are contiguous in GPU memory. Hence, we must call it again
+        self._flatten_parameters()
 
         self.step += 1
         bsize = x.size(0)
@@ -226,14 +241,14 @@ def generate(self, mels, save_path, batched, target, overlap, mu_law):
         output = output.cpu().numpy()
         output = output.astype(np.float64)
 
+        if mu_law:
+            output = decode_mu_law(output, self.n_classes, False)
+
         if batched:
             output = self.xfade_and_unfold(output, target, overlap)
         else:
             output = output[0]
 
-        if mu_law:
-            output = decode_mu_law(output, self.n_classes, False)
-
         # Fade-out at the end to avoid signal cutting out suddenly
         fade_out = np.linspace(1, 0, 20 * self.hop_length)
         output = output[:wave_len]
@@ -388,9 +403,12 @@ def xfade_and_unfold(self, y, target, overlap):
     def get_step(self):
         return self.step.data.item()
 
-    def checkpoint(self, path):
+    def checkpoint(self, path, optimizer):
+        # Optimizer can be given as an argument because checkpoint function is
+        # only useful in context of already existing training process.
         k_steps = self.get_step() // 1000
         self.save(f'{path}/checkpoint_{k_steps}k_steps.pyt')
+        torch.save(optimizer.get_state(), f'{path}/checkpoint_{k_steps}k_steps_optim.pyt')
 
     def log(self, path, msg):
         with open(path, 'a') as f:
@@ -404,15 +422,25 @@ def restore(self, path):
             print(f'\nLoading Weights: "{path}"\n')
             self.load(path)
 
-    def load(self, path, device='cpu'):
-        # because PyTorch places on CPU by default, we follow those semantics by using CPU as default.
+    def load(self, path):
+        # Use device of model params as location for loaded state
+        device = next(self.parameters()).device
         self.load_state_dict(torch.load(path, map_location=device), strict=False)
 
     def save(self, path):
+        # No optimizer argument because saving a model should not include data
+        # only relevant in the training process - it should only be properties
+        # of the model itself. Let caller take care of saving optimzier state.
         torch.save(self.state_dict(), path)
 
     def num_params(self, print_out=True):
         parameters = filter(lambda p: p.requires_grad, self.parameters())
         parameters = sum([np.prod(p.size()) for p in parameters]) / 1_000_000
         if print_out:
             print('Trainable Parameters: %.3fM' % parameters)
+        return parameters
+
+    def _flatten_parameters(self):
+        """Calls `flatten_parameters` on all the rnns used by the WaveRNN. Used
+        to improve efficiency and avoid PyTorch yelling at us."""
+        [m.flatten_parameters() for m in self._to_flatten]

models/tacotron.py

@@ -54,6 +54,9 @@ class CBHG(nn.Module):
     def __init__(self, K, in_channels, channels, proj_channels, num_highways):
         super().__init__()
 
+        # List of all rnns to call `flatten_parameters()` on
+        self._to_flatten = []
+        
         self.bank_kernels = [i for i in range(1, K + 1)]
         self.conv1d_bank = nn.ModuleList()
         for k in self.bank_kernels:
@@ -78,8 +81,16 @@ def __init__(self, K, in_channels, channels, proj_channels, num_highways):
             self.highways.append(hn)
 
         self.rnn = nn.GRU(channels, channels, batch_first=True, bidirectional=True)
+        self._to_flatten.append(self.rnn)
+
+        # Avoid fragmentation of RNN parameters and associated warning
+        self._flatten_parameters()
 
     def forward(self, x):
+        # Although we `_flatten_parameters()` on init, when using DataParallel 
+        # the model gets replicated, making it no longer guaranteed that the
+        # weights are contiguous in GPU memory. Hence, we must call it again
+        self._flatten_parameters()
 
         # Save these for later
         residual = x
@@ -114,6 +125,10 @@ def forward(self, x):
         x, _ = self.rnn(x)
         return x
 
+    def _flatten_parameters(self):
+        """Calls `flatten_parameters` on all the rnns used by the WaveRNN. Used
+        to improve efficiency and avoid PyTorch yelling at us."""
+        [m.flatten_parameters() for m in self._to_flatten]
 
 class PreNet(nn.Module):
     def __init__(self, in_dims, fc1_dims=256, fc2_dims=128, dropout=0.5):
@@ -189,10 +204,12 @@ def forward(self, encoder_seq_proj, query, t):
 
 
 class Decoder(nn.Module):
+    # Class variable because its value doesn't change between classes
+    # yet ought to be scoped by class because its a property of a Decoder
+    max_r = 20
     def __init__(self, n_mels, decoder_dims, lstm_dims):
         super().__init__()
-        self.max_r = 20
-        self.r = None
+        self.register_buffer('r', torch.tensor(1, dtype=torch.int))
         self.generating = False
         self.n_mels = n_mels
         self.prenet = PreNet(n_mels)
@@ -204,8 +221,7 @@ def __init__(self, n_mels, decoder_dims, lstm_dims):
         self.mel_proj = nn.Linear(lstm_dims, n_mels * self.max_r, bias=False)
 
     def zoneout(self, prev, current, p=0.1):
-        device = prev.device
-        assert prev.device == current.device
+        device = next(self.parameters()).device  # Use same device as parameters
         mask = torch.zeros(prev.size(), device=device).bernoulli_(p)
         return prev * mask + current * (1 - mask)
 
@@ -279,17 +295,15 @@ def __init__(self, embed_dims, num_chars, encoder_dims, decoder_dims, n_mels, ff
         self.init_model()
         self.num_params()
 
-        # Unfortunately I have to put these settings into params in order to save
-        # if anyone knows a better way of doing this please open an issue in the repo
-        self.step = nn.Parameter(torch.zeros(1).long(), requires_grad=False)
-        self.r = nn.Parameter(torch.tensor(0).long(), requires_grad=False)
-
-    def set_r(self, r):
-        self.r.data = torch.tensor(r)
-        self.decoder.r = r
+        self.register_buffer('step', torch.zeros(1, dtype=torch.long))
+    
+    @property
+    def r(self):
+        return self.decoder.r.item()
 
-    def get_r(self):
-        return self.r.item()
+    @r.setter
+    def r(self, value):
+        self.decoder.r = self.decoder.r.new_tensor(value, requires_grad=False)
 
     def forward(self, x, m, generate_gta=False):
         device = next(self.parameters()).device  # use same device as parameters
@@ -351,7 +365,7 @@ def forward(self, x, m, generate_gta=False):
 
         # For easy visualisation
         attn_scores = torch.cat(attn_scores, 1)
-        attn_scores = attn_scores.cpu().data.numpy()
+        # attn_scores = attn_scores.cpu().data.numpy()
 
         return mel_outputs, linear, attn_scores
 
@@ -430,11 +444,17 @@ def get_step(self):
         return self.step.data.item()
 
     def reset_step(self):
-        self.step = nn.Parameter(torch.zeros(1).long(), requires_grad=False)
+        assert self.step is not None
+        device = next(self.parameters()).device  # use same device as parameters
+        # assignment to parameters or buffers is overloaded, updates internal dict entry
+        self.step = torch.zeros(1, dtype=torch.long, device=device)
 
-    def checkpoint(self, path):
+    def checkpoint(self, path, optimizer):
+        # Optimizer can be given as an argument because checkpoint function is
+        # only useful in context of already existing training process.
         k_steps = self.get_step() // 1000
         self.save(f'{path}/checkpoint_{k_steps}k_steps.pyt')
+        torch.save(optimizer.get_state(), f'{path}/checkpoint_{k_steps}k_steps_optim.pyt')
 
     def log(self, path, msg):
         with open(path, 'a') as f:
@@ -447,17 +467,21 @@ def restore(self, path):
         else:
             print(f'\nLoading Weights: "{path}"\n')
             self.load(path)
-            self.decoder.r = self.r.item()
 
-    def load(self, path, device='cpu'):
-        # because PyTorch places on CPU by default, we follow those semantics by using CPU as default.
+    def load(self, path):
+        # Use device of model params as location for loaded state
+        device = next(self.parameters()).device
         self.load_state_dict(torch.load(path, map_location=device), strict=False)
 
     def save(self, path):
+        # No optimizer argument because saving a model should not include data
+        # only relevant in the training process - it should only be properties
+        # of the model itself. Let caller take care of saving optimzier state.
         torch.save(self.state_dict(), path)
 
     def num_params(self, print_out=True):
         parameters = filter(lambda p: p.requires_grad, self.parameters())
         parameters = sum([np.prod(p.size()) for p in parameters]) / 1_000_000
         if print_out:
             print('Trainable Parameters: %.3fM' % parameters)
+        return parameters

preprocess.py

@@ -10,9 +10,17 @@
 from utils.files import get_files
 
 
+# Helper functions for argument types
+def valid_n_workers(num):
+    n = int(num)
+    if n < 1:
+        raise argparse.ArgumentTypeError('%r must be an integer greater than 0' % num)
+    return n
+
 parser = argparse.ArgumentParser(description='Preprocessing for WaveRNN and Tacotron')
 parser.add_argument('--path', '-p', default=hp.wav_path, help='directly point to dataset path (overrides hparams.wav_path')
-parser.add_argument('--extension', '-e', default='.wav', help='file extension to search for in dataset folder')
+parser.add_argument('--extension', '-e', metavar='EXT', default='.wav', help='file extension to search for in dataset folder')
+parser.add_argument('--num_workers', '-w', metavar='EXT', type=valid_n_workers, default=cpu_count()-1, help='The number of worker threads to use for preprocessing')
 args = parser.parse_args()
 
 extension = args.extension
@@ -60,13 +68,17 @@ def process_wav(path):
         with open(f'{paths.data}text_dict.pkl', 'wb') as f:
             pickle.dump(text_dict, f)
 
-    simple_table([('Sample Rate', hp.sample_rate),
-                  ('Bit Depth', hp.bits),
-                  ('Mu Law', hp.mu_law),
-                  ('Hop Length', hp.hop_length),
-                  ('CPU Count', cpu_count())])
+    n_workers = max(1, args.num_workers)
+
+    simple_table([
+        ('Sample Rate', hp.sample_rate),
+        ('Bit Depth', hp.bits),
+        ('Mu Law', hp.mu_law),
+        ('Hop Length', hp.hop_length),
+        ('CPU Usage', f'{n_workers}/{cpu_count()}')
+    ])
 
-    pool = Pool(processes=cpu_count())
+    pool = Pool(processes=n_workers)
     dataset = []
 
     for i, (id, length) in enumerate(pool.imap_unordered(process_wav, wav_files), 1):

quick_start.py

@@ -119,6 +119,6 @@
         m = torch.tensor(m).unsqueeze(0)
         m = (m + 4) / 8
 
-        voc_model.generate(m, save_path, batched, hp.voc_target, hp.voc_overlap, hp.mu_law)
+        voc_model.generate(m, save_path, batched, target, overlap, hp.mu_law)
 
     print('\n\nDone.\n')