Tacotron2: `CUDA out of memory` at torch.load(waveglow_path)

[Konard]

Trying to start a demo as it is written in README.md and error occurs (not enough GPU memory).

Note: I have Geforce GTX 950 with 2 GB of GPU memory.

Is it possible to run a demo using this video adapter?

Videos of attempt:
https://youtu.be/Y8xzLLtYy70 (part 1)
https://youtu.be/94R22URpg88 (part 2)
https://youtu.be/6_7vrx6MGBY (part 3)

Stacktrace:

RuntimeError                              Traceback (most recent call last)
<ipython-input-5-7a43ba882a5e> in <module>
      1 waveglow_path = 'waveglow_old.pt'
----> 2 waveglow = torch.load(waveglow_path)['model']
      3 waveglow.cuda()
      4 denoiser = Denoiser(waveglow)

C:\ProgramData\Anaconda3\lib\site-packages\torch\serialization.py in load(f, map_location, pickle_module)
    366         f = open(f, 'rb')
    367     try:
--> 368         return _load(f, map_location, pickle_module)
    369     finally:
    370         if new_fd:

C:\ProgramData\Anaconda3\lib\site-packages\torch\serialization.py in _load(f, map_location, pickle_module)
    540     unpickler = pickle_module.Unpickler(f)
    541     unpickler.persistent_load = persistent_load
--> 542     result = unpickler.load()
    543 
    544     deserialized_storage_keys = pickle_module.load(f)

C:\ProgramData\Anaconda3\lib\site-packages\torch\serialization.py in persistent_load(saved_id)
    503             if root_key not in deserialized_objects:
    504                 deserialized_objects[root_key] = restore_location(
--> 505                     data_type(size), location)
    506             storage = deserialized_objects[root_key]
    507             if view_metadata is not None:

C:\ProgramData\Anaconda3\lib\site-packages\torch\serialization.py in default_restore_location(storage, location)
    112 def default_restore_location(storage, location):
    113     for _, _, fn in _package_registry:
--> 114         result = fn(storage, location)
    115         if result is not None:
    116             return result

C:\ProgramData\Anaconda3\lib\site-packages\torch\serialization.py in _cuda_deserialize(obj, location)
     94     if location.startswith('cuda'):
     95         device = validate_cuda_device(location)
---> 96         return obj.cuda(device)
     97 
     98 

C:\ProgramData\Anaconda3\lib\site-packages\torch\_utils.py in _cuda(self, device, non_blocking, **kwargs)
     74         else:
     75             new_type = getattr(torch.cuda, self.__class__.__name__)
---> 76             return new_type(self.size()).copy_(self, non_blocking)
     77 
     78 

RuntimeError: CUDA out of memory. Tried to allocate 2.00 MiB (GPU 0; 2.00 GiB total capacity; 1.18 GiB already allocated; 1.64 MiB free; 3.18 MiB cached)

[rafaelvalle]

Try doing it in two steps:
1) Load only Tacotron 2 and save it outputs to disk.
2) Load only WaveGlow and load previous taco2 outputs from disk.

[Konard]

I'm not familiar with python, so it will take some time to figure out how to do both your steps. But I've tried just to run Load WaveGlow for mel2audio synthesis and denoiser step, without running Load model from checkpoint step, and I end up with _almost_ the same result, all memory goes out, but stack trace and error change.

RuntimeError                              Traceback (most recent call last)
<ipython-input-4-7a43ba882a5e> in <module>
      2 waveglow = torch.load(waveglow_path)['model']
      3 waveglow.cuda()
----> 4 denoiser = Denoiser(waveglow)

~\Desktop\tacotron2\denoiser.py in __init__(self, waveglow, filter_length, n_overlap, win_length, mode)
     28 
     29         with torch.no_grad():
---> 30             bias_audio = waveglow.infer(mel_input, sigma=0.0).float()
     31             bias_spec, _ = self.stft.transform(bias_audio)
     32 

~\Desktop\tacotron2\waveglow\glow_old.py in infer(self, spect, sigma)
    208                 audio = torch.cat([audio_1, audio[:,n_half:,:]], 1)
    209 
--> 210             audio = self.convinv[k](audio, reverse=True)
    211 
    212             if k%4 == 0 and k > 0:

C:\ProgramData\Anaconda3\lib\site-packages\torch\nn\modules\module.py in __call__(self, *input, **kwargs)
    487             result = self._slow_forward(*input, **kwargs)
    488         else:
--> 489             result = self.forward(*input, **kwargs)
    490         for hook in self._forward_hooks.values():
    491             hook_result = hook(self, input, result)

~\Desktop\tacotron2\waveglow\glow.py in forward(self, z, reverse)
     94                     W_inverse = W_inverse.half()
     95                 self.W_inverse = W_inverse
---> 96             z = F.conv1d(z, self.W_inverse, bias=None, stride=1, padding=0)
     97             return z
     98         else:

RuntimeError: CUDA error: out of memory

So if I change Load WaveGlow for mel2audio synthesis and denoiser task to:

waveglow_path = 'waveglow_old.pt'
waveglow = torch.load(waveglow_path)['model']
waveglow.cuda()
#denoiser = Denoiser(waveglow)

commenting out the Denoiser completes the loading.

Looks like WaveGlow loading result should be dumped to _disk?_
Wait, I have 64 GB of RAM on my machine, is it possible just to free Cuda memory, but store intermediate results at RAM? Using the disk should be a lot slower, also I don`t have too less of disk space left for now.

If I only running Load model from checkpoint step it fills up 1 GB of my GPU memory.

[Konard]

Ok, I think I've figured out how to do it:

Load model from checkpoint

checkpoint_path = "tacotron2_statedict.pt"
model = load_model(hparams)
model.load_state_dict(torch.load(checkpoint_path)['state_dict'])
_ = model.eval()

statedictBuffer = io.BytesIO()
torch.save(model, statedictBuffer)
del model;
torch.cuda.empty_cache()

Load WaveGlow for mel2audio synthesis and denoiser

waveglow_path = 'waveglow_old.pt'
waveglow = torch.load(waveglow_path)['model']
waveglow.cuda()

waveglowBuffer = io.BytesIO()
torch.save(waveglow, waveglowBuffer)
del waveglow;
torch.cuda.empty_cache()

waveglowBuffer.seek(0)
waveglow = torch.load(waveglowBuffer)

denoiser = Denoiser(waveglow)

But this also ends up with out of memory error.

RuntimeError                              Traceback (most recent call last)
<ipython-input-5-d4516447ef56> in <module>
     11 waveglow = torch.load(waveglowBuffer)
     12 
---> 13 denoiser = Denoiser(waveglow)

~\Desktop\tacotron2\denoiser.py in __init__(self, waveglow, filter_length, n_overlap, win_length, mode)
     28 
     29         with torch.no_grad():
---> 30             bias_audio = waveglow.infer(mel_input, sigma=0.0).float()
     31             bias_spec, _ = self.stft.transform(bias_audio)
     32 

~\Desktop\tacotron2\waveglow\glow_old.py in infer(self, spect, sigma)
    199                 audio_0 = audio[:,n_half:,:]
    200 
--> 201             output = self.WN[k]((audio_0, spect))
    202             s = output[:, n_half:, :]
    203             b = output[:, :n_half, :]

C:\ProgramData\Anaconda3\lib\site-packages\torch\nn\modules\module.py in __call__(self, *input, **kwargs)
    487             result = self._slow_forward(*input, **kwargs)
    488         else:
--> 489             result = self.forward(*input, **kwargs)
    490         for hook in self._forward_hooks.values():
    491             hook_result = hook(self, input, result)

~\Desktop\tacotron2\waveglow\glow_old.py in forward(self, forward_input)
     70             acts = fused_add_tanh_sigmoid_multiply(
     71                 self.in_layers[i](audio),
---> 72                 self.cond_layers[i](spect),
     73                 torch.IntTensor([self.n_channels]))
     74 

C:\ProgramData\Anaconda3\lib\site-packages\torch\nn\modules\module.py in __call__(self, *input, **kwargs)
    487             result = self._slow_forward(*input, **kwargs)
    488         else:
--> 489             result = self.forward(*input, **kwargs)
    490         for hook in self._forward_hooks.values():
    491             hook_result = hook(self, input, result)

C:\ProgramData\Anaconda3\lib\site-packages\torch\nn\modules\conv.py in forward(self, input)
    185     def forward(self, input):
    186         return F.conv1d(input, self.weight, self.bias, self.stride,
--> 187                         self.padding, self.dilation, self.groups)
    188 
    189 

RuntimeError: CUDA out of memory. Tried to allocate 6.88 MiB (GPU 0; 2.00 GiB total capacity; 1.32 GiB already allocated; 1.94 MiB free; 1.69 MiB cached)

Is this actually a correct way to free CUDA memory?

del model;
torch.cuda.empty_cache()

Looks like this frees only a few megabytes of GPU memory.

[Konard]

Now I'm using kernel restart in jupyter notebook to free GPU memory.
Also, I actually tried to use the disk as intermediate storage.
So the code snippets are like this:

Load model from checkpoint (restored)

checkpoint_path = "tacotron2_statedict.pt"
model = load_model(hparams)
model.load_state_dict(torch.load(checkpoint_path)['state_dict'])
_ = model.eval()

Load WaveGlow for mel2audio synthesis and denoiser (no denoiser)

waveglow_path = 'waveglow_old.pt'
waveglow = torch.load(waveglow_path)['model']
waveglow.cuda()

torch.save(waveglow, 'waveglow.tmp')

I run this separately doing dump on disk and restart of the kernel to free GPU memory.

Decode text input and plot results (now dumps to disk)

mel_outputs, mel_outputs_postnet, _, alignments = model.inference(sequence)
plot_data((mel_outputs.data.cpu().numpy()[0],
           mel_outputs_postnet.data.cpu().numpy()[0],
           alignments.data.cpu().numpy()[0].T))
torch.save(mel_outputs_postnet, 'tensor.tmp')

I run this also seperately doing dump on disk and restart of kernel to free GPU memory.

When everything ready on disk, and GPU memory is free, I run only steps required for Synthesize:

Synthesize audio from spectrogram using WaveGlow (loading both models from disk to GPU)

waveglow = torch.load('waveglow.tmp')
mel_outputs_postnet = torch.load('tensor.tmp')
with torch.no_grad():
    audio = waveglow.infer(mel_outputs_postnet, sigma=0.666)
ipd.Audio(audio[0].data.cpu().numpy(), rate=hparams.sampling_rate)

And the result is the same on GPU (out of memory error):

RuntimeError                              Traceback (most recent call last)
<ipython-input-4-8ecc19d9c3f7> in <module>
      5 mel_outputs_postnet = torch.load('tensor.tmp')
      6 with torch.no_grad():
----> 7     audio = waveglow.infer(mel_outputs_postnet, sigma=0.666)
      8 ipd.Audio(audio[0].data.cpu().numpy(), rate=hparams.sampling_rate)

~\Desktop\tacotron2\waveglow\glow_old.py in infer(self, spect, sigma)
    199                 audio_0 = audio[:,n_half:,:]
    200 
--> 201             output = self.WN[k]((audio_0, spect))
    202             s = output[:, n_half:, :]
    203             b = output[:, :n_half, :]

C:\ProgramData\Anaconda3\lib\site-packages\torch\nn\modules\module.py in __call__(self, *input, **kwargs)
    487             result = self._slow_forward(*input, **kwargs)
    488         else:
--> 489             result = self.forward(*input, **kwargs)
    490         for hook in self._forward_hooks.values():
    491             hook_result = hook(self, input, result)

~\Desktop\tacotron2\waveglow\glow_old.py in forward(self, forward_input)
     69         for i in range(self.n_layers):
     70             acts = fused_add_tanh_sigmoid_multiply(
---> 71                 self.in_layers[i](audio),
     72                 self.cond_layers[i](spect),
     73                 torch.IntTensor([self.n_channels]))

C:\ProgramData\Anaconda3\lib\site-packages\torch\nn\modules\module.py in __call__(self, *input, **kwargs)
    487             result = self._slow_forward(*input, **kwargs)
    488         else:
--> 489             result = self.forward(*input, **kwargs)
    490         for hook in self._forward_hooks.values():
    491             hook_result = hook(self, input, result)

C:\ProgramData\Anaconda3\lib\site-packages\torch\nn\modules\conv.py in forward(self, input)
    185     def forward(self, input):
    186         return F.conv1d(input, self.weight, self.bias, self.stride,
--> 187                         self.padding, self.dilation, self.groups)
    188 
    189 

RuntimeError: CUDA out of memory. Tried to allocate 125.00 MiB (GPU 0; 2.00 GiB total capacity; 1.32 GiB already allocated; 15.46 MiB free; 10.78 MiB cached)

But it is actually possible to run this on CPU:

Synthesize audio from spectrogram using WaveGlow (loading both models from disk to CPU)

device = torch.device('cpu')
waveglow = torch.load('waveglow.tmp', map_location=device)
mel_outputs_postnet = torch.load('tensor.tmp', map_location=device)
with torch.no_grad():
    audio = waveglow.infer(mel_outputs_postnet, sigma=0.666)
ipd.Audio(audio[0].data.cpu().numpy(), rate=hparams.sampling_rate)

In order this to work I had to fix tacotron2\waveglow\glow_old.py file. On lines 186 and 218 it is required to replace torch.cuda.FloatTensor with torch.FloatTensor.

And this is finally working, but slow. Btw infer loads 4 out of 8 logical CPUs and allocates on average about 3 GB of RAM.

So I think there should be a way to work with waveglow that is smaller in size or to use it in chunks.

To apply a Denoiser on CPU, I also had to fix it, removing .cuda() at 15 and 36 lines of tacotron2\denoiser.py. tacotron2\denoiser.py itself was copied to my working directory from here: https://github.com/NVIDIA/waveglow/blob/61adc104147f569b9136c39bb8ad296f6d42bb43/denoiser.py

(Optional) Remove WaveGlow bias (working on CPU)

denoiser = Denoiser(waveglow)

audio_denoised = denoiser(audio, strength=0.01)[:, 0]
ipd.Audio(audio_denoised.cpu().numpy(), rate=hparams.sampling_rate) 

Generated audio files:
audio.zip

Video of CPU speech synthesis:
https://youtu.be/VmHcjX2eAQA

Forks with all required changes:
https://github.com/Konard/waveglow
https://github.com/Konard/tacotron2

[loretoparisi]

@rafaelvalle it seems we are having a similar issue when inferencing with the new model on CPU https://github.com/NVIDIA/tacotron2/issues/192