126 lines
4.9 KiB
Python
126 lines
4.9 KiB
Python
# Original source: https://github.com/SeanNaren/deepspeech.pytorch/blob/master/deepspeech_pytorch/loader/sparse_image_warp.py
|
|
# Removes the time_warp augmentation and only implements masking.
|
|
|
|
import numpy as np
|
|
import random
|
|
|
|
import torch
|
|
import torchvision.utils
|
|
|
|
from trainer.inject import Injector
|
|
from utils.util import opt_get
|
|
|
|
|
|
def spec_augment(mel_spectrogram, frequency_masking_para=27, time_masking_para=70, frequency_mask_num=1, time_mask_num=1):
|
|
|
|
v = mel_spectrogram.shape[1]
|
|
tau = mel_spectrogram.shape[2]
|
|
|
|
# Step 2 : Frequency masking
|
|
for i in range(frequency_mask_num):
|
|
f = np.random.uniform(low=0.0, high=frequency_masking_para)
|
|
f = int(f)
|
|
if v - f < 0:
|
|
continue
|
|
f0 = random.randint(0, v-f)
|
|
mel_spectrogram[:, f0:f0+f, :] = 0
|
|
|
|
# Step 3 : Time masking
|
|
for i in range(time_mask_num):
|
|
t = np.random.uniform(low=0.0, high=time_masking_para)
|
|
t = int(t)
|
|
if tau - t < 0:
|
|
continue
|
|
t0 = random.randint(0, tau-t)
|
|
mel_spectrogram[:, :, t0:t0+t] = 0
|
|
|
|
return mel_spectrogram
|
|
|
|
|
|
class MelMaskInjector(Injector):
|
|
def __init__(self, opt, env):
|
|
super().__init__(opt, env)
|
|
self.freq_mask_sz = opt_get(opt, ['frequency_mask_size_high'], 27)
|
|
self.n_freq_masks = opt_get(opt, ['frequency_mask_count'], 1)
|
|
self.time_mask_sz = opt_get(opt, ['time_mask_size_high'], 5)
|
|
self.n_time_masks = opt_get(opt, ['time_mask_count'], 3)
|
|
|
|
def forward(self, state):
|
|
h = state[self.input]
|
|
return {self.output: spec_augment(h, self.freq_mask_sz, self.time_mask_sz, self.n_freq_masks, self.n_time_masks)}
|
|
|
|
def visualization_spectrogram(spec, title):
|
|
# Turns spec into an image and outputs it to the filesystem.
|
|
spec = spec.unsqueeze(dim=1)
|
|
# Normalize so spectrogram is easier to view.
|
|
spec = (spec - spec.mean()) / spec.std()
|
|
spec = ((spec + 1) / 2).clip(0, 1)
|
|
torchvision.utils.save_image(spec, f'{title}.png')
|
|
|
|
|
|
def test_mel_mask():
|
|
from data.audio.unsupervised_audio_dataset import load_audio
|
|
from trainer.injectors.base_injectors import MelSpectrogramInjector
|
|
spec_maker = MelSpectrogramInjector({'in': 'audio', 'out': 'spec'}, {})
|
|
a = load_audio('D:\\data\\audio\\libritts\\test-clean\\61\\70970\\61_70970_000007_000001.wav', 22050).unsqueeze(0)
|
|
s = spec_maker({'audio': a})['spec']
|
|
visualization_spectrogram(s, 'original spec')
|
|
saug = spec_augment(s, 50, 5, 1, 3)
|
|
visualization_spectrogram(saug, 'modified spec')
|
|
|
|
|
|
'''
|
|
Crafty bespoke injector that is used when training ASR models to create longer sequences to ensure that the entire
|
|
input length embedding is trained. Does this by concatenating every other batch element together to create longer
|
|
sequences which (theoretically) use similar amounts of GPU memory.
|
|
'''
|
|
class CombineMelInjector(Injector):
|
|
def __init__(self, opt, env):
|
|
super().__init__(opt, env)
|
|
self.audio_key = opt['audio_key']
|
|
self.text_key = opt['text_key']
|
|
self.audio_lengths = opt['audio_lengths_key']
|
|
self.text_lengths = opt['text_lengths_key']
|
|
self.output_audio_key = opt['output_audio_key']
|
|
self.output_text_key = opt['output_text_key']
|
|
from models.tacotron2.text import symbols
|
|
self.text_separator = len(symbols)+1 # Probably need to allow this to be set by user.
|
|
|
|
def forward(self, state):
|
|
audio = state[self.audio_key]
|
|
texts = state[self.text_key]
|
|
audio_lengths = state[self.audio_lengths]
|
|
text_lengths = state[self.text_lengths]
|
|
if audio.shape[0] == 1:
|
|
return {self.output_audio_key: audio, self.output_text_key: texts}
|
|
combined_audios = []
|
|
combined_texts = []
|
|
for b in range(audio.shape[0]//2):
|
|
a1 = audio[b*2, :, :audio_lengths[b*2]]
|
|
a2 = audio[b*2+1, :, :audio_lengths[b*2+1]]
|
|
a = torch.cat([a1, a2], dim=1)
|
|
a = torch.nn.functional.pad(a, (0, audio.shape[-1]*2-a.shape[-1]))
|
|
combined_audios.append(a)
|
|
|
|
t1 = texts[b*2, :text_lengths[b*2]]
|
|
t1 = torch.nn.functional.pad(t1, (0, 1), value=self.text_separator)
|
|
t2 = texts[b*2+1, :text_lengths[b*2+1]]
|
|
t = torch.cat([t1, t2], dim=0)
|
|
t = torch.nn.functional.pad(t, (0, texts.shape[-1]*2-t.shape[-1]))
|
|
combined_texts.append(t)
|
|
return {self.output_audio_key: torch.stack(combined_audios, dim=0),
|
|
self.output_text_key: torch.stack(combined_texts, dim=0)}
|
|
|
|
|
|
def test_mel_injector():
|
|
inj = CombineMelInjector({'audio_key': 'a', 'text_key': 't', 'audio_lengths_key': "alk", 'text_lengths_key': 'tlk'}, {})
|
|
a = torch.rand((4, 22000))
|
|
al = torch.tensor([11000,14000,22000,20000])
|
|
t = torch.randint(0, 120, (4, 250))
|
|
tl = torch.tensor([100,120,200,250])
|
|
rs = inj({'a': a, 't': t, 'alk': al, 'tlk': tl})
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
test_mel_injector() |