Building an End-to-End Speech Recognition Model in PyTorch with AssemblyAI

This post was written by Michael Nguyen, Machine Learning Research Engineer at AssemblyAI . AssemblyAI uses Comet to log, visualize, and understand their model development pipeline. Deep Learning has changed the game in speech recognition with the introduction of end-to-end models. These models take in audio, and directly output transcriptions. Two of the most popular end-to-end models today are Deep Speech by Baidu, and Listen Attend Spell (LAS) by Google. Both Deep Speech and LAS, are recurrent neural network (RNN) based architectures with different approaches to modeling speech recognition. Deep Speech uses the Connectionist Temporal Classification (CTC) loss function to predict the speech transcript. LAS uses a sequence to sequence network architecture for its predictions. These models simplified speech recognition pipelines by taking advantage of the capacity of deep learning system to learn from large datasets. With enough data, you should, in theory, be able to build a super robust speech recognition model that can account for all the nuance in speech without having to spend a ton of time and effort hand engineering acoustic features or dealing with complex pipelines in more old-school GMM-HMM model architectures, for example. Deep learning is a fast-moving field, and Deep Speech and LAS style architectures are already quickly becoming outdated. You can read about where the industry is moving in the Latest Advancement Section below. How to Build Your Own End-to-End Speech Recognition Model in PyTorch Let’s walk through how one would build their own end-to-end speech recognition model in PyTorch. The model we’ll build is inspired by Deep Speech 2 (Baidu’s second revision of their now-famous model) with some personal improvements to the architecture. The output of the model will be a probability matrix of characters, and we’ll use that probability matrix to decode the most likely characters spoken from the audio. You can find the full code and also run the it with GPU support on . Google Colaboratory Preparing the data pipeline Data is one of the most important aspects of speech recognition. We’ll take raw audio waves and transform them into Mel Spectrograms. You can read more on the details about how that transformation looks from this excellent post . For this post, you can just think of a Mel Spectrogram as essentially a picture of sound. here For handling the audio data, we are going to use an extremely useful utility called which is a library built by the PyTorch team specifically for audio data. We’ll be training on a subset of , which is a corpus of read English speech data derived from audiobooks, comprising 100 hours of transcribed audio data. You can easily download this dataset using : torchaudio LibriSpeech torchaudio torchaudio train_dataset = torchaudio.datasets.LIBRISPEECH( , url= , download=True) test_dataset = torchaudio.datasets.LIBRISPEECH( , url= , download=True) import "./" "train-clean-100" "./" "test-clean" Each sample of the dataset contains the waveform, sample rate of audio, the utterance/label, and more metadata on the sample. You can view what each sample looks like from the source code . here Data Augmentation – SpecAugment Data augmentation is a technique used to artificially increase the diversity of your dataset in order to increase your dataset size. This strategy is especially helpful when data is scarce or if your model is overfitting. For speech recognition, you can do the standard augmentation techniques, like changing the pitch, speed, injecting noise, and adding reverb to your audio data. We found Spectrogram Augmentation (SpecAugment), to be a much simpler and more effective approach. SpecAugment, was first introduced in the paper , in which the authors found that simply cutting out random blocks of consecutive time and frequency dimensions improved the models generalization abilities significantly! SpecAugment: A Simple Data Augmentation Method for Automatic Speech Recognition In PyTorch, you can use the function to mask out the frequency dimension, and for the time dimension. torchaudio FrequencyMasking TimeMasking torchaudio () torchaudio () .transforms .FrequencyMasking .transforms .TimeMasking Now that we have the data, we’ll need to transform the audio into Mel Spectrograms, and map the character labels for each audio sample into integer labels: char_map_str = self.char_map = {} self.index_map = {} line char_map_str.strip().split( ): ch, index = line.split() self.char_map[ch] = int(index) self.index_map[int(index)] = ch self.index_map[ ] = int_sequence = [] c text: c == : ch = self.char_map[ ] : ch = self.char_map[c] int_sequence.append(ch) int_sequence string = [] i labels: string.append(self.index_map[i]) .join(string).replace( , ) train_audio_transforms = nn.Sequential( torchaudio.transforms.MelSpectrogram(sample_rate= , n_mels= ), torchaudio.transforms.FrequencyMasking(freq_mask_param= ), torchaudio.transforms.TimeMasking(time_mask_param= ) ) valid_audio_transforms = torchaudio.transforms.MelSpectrogram() text_transform = TextTransform() spectrograms = [] labels = [] input_lengths = [] label_lengths = [] (waveform, _, utterance, _, _, _) data: data_type == : spec = train_audio_transforms(waveform).squeeze( ).transpose( , ) : spec = valid_audio_transforms(waveform).squeeze( ).transpose( , ) spectrograms.append(spec) label = torch.Tensor(text_transform.text_to_int(utterance.lower())) labels.append(label) input_lengths.append(spec.shape[ ]// ) label_lengths.append(len(label)) spectrograms = nn.utils.rnn.pad_sequence(spectrograms, batch_first= ).unsqueeze( ).transpose( , ) labels = nn.utils.rnn.pad_sequence(labels, batch_first= ) spectrograms, labels, input_lengths, label_lengths : class TextTransform """Maps characters to integers and vice versa""" : def __init__ (self) """ ' 0 &lt;SPACE&gt; 1 a 2 b 3 c 4 d 5 e 6 f 7 g 8 h 9 i 10 j 11 k 12 l 13 m 14 n 15 o 16 p 17 q 18 r 19 s 20 t 21 u 22 v 23 w 24 x 25 y 26 z 27 """ for in '\n' 1 ' ' : def text_to_int (self, text) """ Use a character map and convert text to an integer sequence """ for in if ' ' '' else return : def int_to_text (self, labels) """ Use a character map and convert integer labels to an text sequence """ for in return '' '' ' ' 16000 128 15 35 : def data_processing (data, data_type= ) "train" for in if 'train' 0 0 1 else 0 0 1 0 2 True 1 2 3 True return Define the Model – Deep Speech 2 (but better) Our model will be similar to the Deep Speech 2 architecture. The model will have two main neural network modules – N layers of Residual Convolutional Neural Networks (ResCNN) to learn the relevant audio features, and a set of Bidirectional Recurrent Neural Networks (BiRNN) to leverage the learned ResCNN audio features. The model is topped off with a fully connected layer used to classify characters per time step. Convolutional Neural Networks (CNN) are great at extracting abstract features, and we’ll apply the same feature extraction power to audio spectrograms. Instead of just vanilla CNN layers, we choose to use Residual CNN layers. Residual connections (AKA skip connections) were first introduced in the paper , where the author found that you can build really deep networks with good accuracy gains if you add these connections to your CNN’s. Deep Residual Learning for Image Recognition Adding these Residual connections also helps the model learn faster and generalize better. The paper shows that networks with residual connections have a “flatter” loss surface, making it easier for models to navigate the loss landscape and find a lower and more generalizable minima. Visualizing the Loss Landscape of Neural Nets Recurrent Neural Networks (RNN) are naturally great at sequence modeling problems. RNN’s processes the audio features step by step, making a prediction for each frame while using context from previous frames. We use BiRNN’s because we want the context of not only the frame before each step, but the frames after it as well. This can help the model make better predictions, as each frame in the audio will have more information before making a prediction. We use Gated Recurrent Unit (GRU’s) variant of RNN’s as it needs less computational resources than LSTM’s, and works just as well in some cases. The model outputs a probability matrix for characters which we’ll use to feed into our decoder to extract what the model believes are the highest probability characters that were spoken. ( . ): """ normalization built for cnns input""" def __init__( , ): super( , ).__init__() self.layer_norm = nn. ( ) def forward( , ): # x ( , , , ) x = x.transpose(2, 3).contiguous() # ( , , , ) x = self.layer_norm( ) return x.transpose(2, 3).contiguous() # ( , , , ) ( . ): """ inspired by https://arxiv.org/pdf/1603.05027.pdf except with layer norm instead of batch norm """ def __init__( , , , , , , ): super( , ).__init__() self.cnn1 = nn. ( , , , , = //2) self.cnn2 = nn. ( , , , , = //2) self.dropout1 = nn. ( ) self.dropout2 = nn. ( ) self.layer_norm1 = ( ) self.layer_norm2 = ( ) def forward( , ): residual = x # ( , , , ) x = self.layer_norm1( ) x = .gelu( ) x = self.dropout1( ) x = self.cnn1( ) x = self.layer_norm2( ) x = .gelu( ) x = self.dropout2( ) x = self.cnn2( ) x += residual return x # ( , , , ) ( . ): def __init__( , , , , ): super( , ).__init__() self. = nn. ( = , = , =1, = , = ) self.layer_norm = nn. ( ) self.dropout = nn. ( ) def forward( , ): x = self.layer_norm( ) x = .gelu( ) x, _ = self. ( ) x = self.dropout( ) return x ( . ): """ by 2""" def __init__( , , , , , , =2, =0.1): super( , ).__init__() n_feats = n_feats//2 self.cnn = nn. (1, 32, 3, = , =3//2) # cnn for extracting heirachal features # n residual cnn layers with filter size of 32 self.rescnn_layers = nn. (*[ (32, 32, =3, =1, = , = ) for _ in range( ) ]) self.fully_connected = nn. ( *32, ) self.birnn_layers = nn. (*[ ( = ==0 *2, = , = , = ==0) for i in range( ) ]) self.classifier = nn. ( . ( *2, ), # birnn returns rnn_dim*2 nn. (), nn. ( ), nn. ( , ) ) def forward( , ): x = self.cnn( ) x = self.rescnn_layers( ) sizes = x.size() x = x.view( [0], [1] * [2], [3]) # ( , , ) x = x.transpose(1, 2) # ( , , ) x = self.fully_connected( ) x = self.birnn_layers( ) x = self.classifier( ) return x class CNNLayerNorm nn Module Layer self n_feats CNNLayerNorm self LayerNorm n_feats self x batch channel feature time batch channel time feature x batch channel feature time class ResidualCNN nn Module Residual CNN self in_channels out_channels kernel stride dropout n_feats ResidualCNN self Conv2d in_channels out_channels kernel stride padding kernel Conv2d out_channels out_channels kernel stride padding kernel Dropout dropout Dropout dropout CNNLayerNorm n_feats CNNLayerNorm n_feats self x batch channel feature time x F x x x x F x x x batch channel feature time class BidirectionalGRU nn Module self rnn_dim hidden_size dropout batch_first BidirectionalGRU self BiGRU GRU input_size rnn_dim hidden_size hidden_size num_layers batch_first batch_first bidirectional True LayerNorm rnn_dim Dropout dropout self x x F x BiGRU x x class SpeechRecognitionModel nn Module Speech Recognition Model Inspired DeepSpeech self n_cnn_layers n_rnn_layers rnn_dim n_class n_feats stride dropout SpeechRecognitionModel self Conv2d stride stride padding Sequential ResidualCNN kernel stride dropout dropout n_feats n_feats n_cnn_layers Linear n_feats rnn_dim Sequential BidirectionalGRU rnn_dim rnn_dim if i else rnn_dim hidden_size rnn_dim dropout dropout batch_first i n_rnn_layers Sequential nn Linear rnn_dim rnn_dim GELU Dropout dropout Linear rnn_dim n_class self x x x sizes sizes sizes sizes batch feature time batch time feature x x x Picking the Right Optimizer and Scheduler – AdamW with Super Convergence The optimizer and learning rate schedule plays a very important role in getting our model to converge to the best point. Picking the right optimizer and scheduler can also save you compute time, and help your model generalize better to real-world use cases. For our model, we’ll be using with the . is a widely used optimizer that helps your model converge more quickly, therefore, saving compute time, but has been notorious for not generalizing as well as AKA . AdamW One Cycle Learning Rate Scheduler Adam Stochastic Gradient Descent SGD was first introduced in , and is considered a “fix” to . The paper pointed out that the original algorithm has a wrong implementation of weight decay, which attempts to fix. This fix helps with ‘s generalization problem. AdamW Decoupled Weight Decay Regularization Adam Adam AdamW Adam The was first introduced in the paper . This paper shows that you can train neural networks an order of magnitude faster, while keeping their generalizable abilities, using a simple trick. One Cycle Learning Rate Scheduler Super-Convergence: Very Fast Training of Neural Networks Using Large Learning Rates You start with a low learning rate, which warms up to a large maximum learning rate, then decays linearly to the same point of where you originally started. Because the maximum learning rate is magnitudes higher than the lowest, you also gain some regularization benefits which helps your model generalize better if you have a smaller set of data. With PyTorch, these two methods are already part of the package. optimizer = optim. (model.parameters(), hparams[ ]) scheduler = optim.lr_scheduler. (optimizer, max_lr=hparams[ ], steps_per_epoch=int(len(train_loader)), epochs=hparams[ ], anneal_strategy= ) AdamW 'learning_rate' OneCycleLR 'learning_rate' 'epochs' 'linear' The CTC Loss Function – Aligning Audio to Transcript Our model will be trained to predict the probability distribution of all characters in the alphabet for each frame (ie, timestep) in the spectrogram we feed into the model. Traditional speech recognition models would require you to align the transcript text to the audio before training, and the model would be trained to predict specific labels at specific frames. The innovation of the CTC loss function is that it allows us to skip this step. Our model will learn to align the transcript itself during training. The key to this is the “blank” label introduced by CTC, which gives the model the ability to say that a certain audio frame did not produce a character. You can see a more detailed explanation of CTC and how it works from . this excellent post The CTC loss function is also built into PyTorch. criterion = nn.CTCLoss(blank= ).to(device) 28 Evaluating Your Speech Model When Evaluating your speech recognition model, the industry standard is using the Word Error Rate (WER) as the metric. The Word Error Rate does exactly what it says – it takes the transcription your model outputs, and the true transcription, and measures the error between them. You can see how that’s implemented . Another useful metric is called the Character Error Rate (CER). The CER measures the error of the characters between the model’s output and the true labels. These metrics are helpful to measure how well your model performs. here For this tutorial, we’ll use a “greedy” decoding method to process our model’s output into characters that can be combined to create the transcript. A “greedy” decoder takes in the model output, which is a softmax probability matrix of characters, and for each time step (spectrogram frame), it chooses the label with the highest probability. If the label is a blank label, we remove it from the final transcript. def : arg_maxes = torch.argmax(output, dim= ) decodes = targets = for i, args enumerate(arg_maxes): decode = targets.append(text_transform. )) for j, index enumerate(args): index != blank_label: collapse_repeated j != index args : continue decode.append(index.item ) decodes.append(text_transform. ) return decodes, targets GreedyDecoder( , , , =28, =True) output labels label_lengths blank_label collapse_repeated 2 [] [] in [] int _to_text( [ ][: [ ]]. () labels i label_lengths i tolist in if if and 0 and == [ - ] j 1 () int _to_text( ) decode Training and Monitoring Your Experiments Using Comet.ml provides a platform that allows deep learning researchers to track, compare, explain, and optimize their experiments and models. Comet.ml has improved our productivity at AssemblyAI and we highly recommend using this platform for teams doing any sort of data science experiments. Comet.ml Comet.ml is super easy to set up. And works with just a few lines of code. # initialize experiment object experiment = experiment.set # track metrics experiment.log ) Experiment( = , = ) api_key comet_api_key project_name project_name _name( ) exp_name _metric(' ', . () loss loss item provides you with a very productive dashboard where you can view and track your model’s progress. Comet.ml You can use Comet to track metrics, code, hyper parameters, your model’s graphs, among many other things! A really handy feature that Comet provides is the ability to compare your experiment among many other experiments. Comet has a rich feature set that we won’t cover all here, but we highly recommended using it for a productivity and sanity boost. Here is the rest of our training script. self.val = self.val += self.val model.train() data_len = len(train_loader.dataset) experiment.train(): batch_idx, _data enumerate(train_loader): spectrograms, labels, input_lengths, label_lengths = _data spectrograms, labels = spectrograms.to(device), labels.to(device) optimizer.zero_grad() output = model(spectrograms) output = F.log_softmax(output, dim= ) output = output.transpose( , ) loss = criterion(output, labels, input_lengths, label_lengths) loss.backward() experiment.log_metric( , loss.item(), step=iter_meter.get()) experiment.log_metric( , scheduler.get_lr(), step=iter_meter.get()) optimizer.step() scheduler.step() iter_meter.step() batch_idx % == batch_idx == data_len: print( .format( epoch, batch_idx * len(spectrograms), data_len, * batch_idx / len(train_loader), loss.item())) print( ) model.eval() test_loss = test_cer, test_wer = [], [] experiment.test(): torch.no_grad(): I, _data enumerate(test_loader): spectrograms, labels, input_lengths, label_lengths = _data spectrograms, labels = spectrograms.to(device), labels.to(device) output = model(spectrograms) output = F.log_softmax(output, dim= ) output = output.transpose( , ) loss = criterion(output, labels, input_lengths, label_lengths) test_loss += loss.item() / len(test_loader) decoded_preds, decoded_targets = GreedyDecoder(output.transpose( , ), labels, label_lengths) j range(len(decoded_preds)): test_cer.append(cer(decoded_targets[j], decoded_preds[j])) test_wer.append(wer(decoded_targets[j], decoded_preds[j])) avg_cer = sum(test_cer)/len(test_cer) avg_wer = sum(test_wer)/len(test_wer) experiment.log_metric( , test_loss, step=iter_meter.get()) experiment.log_metric( , avg_cer, step=iter_meter.get()) experiment.log_metric( , avg_wer, step=iter_meter.get()) print( .format(test_loss, avg_cer, avg_wer)) hparams = { : , : , : , : , : , : , : , : learning_rate, : batch_size, : epochs } experiment.log_parameters(hparams) use_cuda = torch.cuda.is_available() torch.manual_seed( ) device = torch.device( use_cuda ) os.path.isdir( ): os.makedirs( ) train_dataset = torchaudio.datasets.LIBRISPEECH( , url=train_url, download= ) test_dataset = torchaudio.datasets.LIBRISPEECH( , url=test_url, download= ) kwargs = { : , : } use_cuda {} train_loader = data.DataLoader(dataset=train_dataset, batch_size=hparams[ ], shuffle= , collate_fn= x: data_processing(x, ), **kwargs) test_loader = data.DataLoader(dataset=test_dataset, batch_size=hparams[ ], shuffle= , collate_fn= x: data_processing(x, ), **kwargs) model = SpeechRecognitionModel( hparams[ ], hparams[ ], hparams[ ], hparams[ ], hparams[ ], hparams[ ], hparams[ ] ).to(device) print(model) print( , sum([param.nelement() param model.parameters()])) optimizer = optim.AdamW(model.parameters(), hparams[ ]) criterion = nn.CTCLoss(blank= ).to(device) scheduler = optim.lr_scheduler.OneCycleLR(optimizer, max_lr=hparams[ ], steps_per_epoch=int(len(train_loader)), epochs=hparams[ ], anneal_strategy= ) iter_meter = IterMeter() epoch range( , epochs + ): train(model, device, train_loader, criterion, optimizer, scheduler, epoch, iter_meter, experiment) test(model, device, test_loader, criterion, epoch, iter_meter, experiment) : class IterMeter (object) """keeps track of total iterations""" : def __init__ (self) 0 : def step (self) 1 : def get (self) return : def train (model, device, train_loader, criterion, optimizer, scheduler, epoch, iter_meter, experiment) with for in # (batch, time, n_class) 2 0 1 # (time, batch, n_class) 'loss' 'learning_rate' if 100 0 or 'Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}' 100. : def test (model, device, test_loader, criterion, epoch, iter_meter, experiment) '\nevaluating…' 0 with with for in # (batch, time, n_class) 2 0 1 # (time, batch, n_class) 0 1 for in 'test_loss' 'cer' 'wer' 'Test set: Average loss: {:.4f}, Average CER: {:4f} Average WER: {:.4f}\n' : def main (learning_rate= , batch_size= , epochs= , train_url= , test_url= , experiment=Experiment ) 5e-4 20 10 "train-clean-100" "test-clean" (api_key= , disabled=True) 'dummy_key' "n_cnn_layers" 3 "n_rnn_layers" 5 "rnn_dim" 512 "n_class" 29 "n_feats" 128 "stride" 2 "dropout" 0.1 "learning_rate" "batch_size" "epochs" 7 "cuda" if else "cpu" if not "./data" "./data" "./data" True "./data" True 'num_workers' 1 'pin_memory' True if else 'batch_size' True lambda 'train' 'batch_size' False lambda 'valid' 'n_cnn_layers' 'n_rnn_layers' 'rnn_dim' 'n_class' 'n_feats' 'stride' 'dropout' 'Num Model Parameters' for in 'learning_rate' 28 'learning_rate' 'epochs' 'linear' for in 1 1 The function trains the model on a full epoch of data. The function evaluates the model on test data after every epoch. It gets the as well as the and of the model. You can start running the training script right now with GPU support in the . train test test_loss cer wer Google Colaboratory How to Improve Accuracy Speech Recognition Requires a ton of data and a ton of compute resources. The example laid out is trained on a subset of LibriSpeech (100 hours of audio) and a single GPU. To get state of the art results you’ll need to do distributed training on thousands of hours of data, on tens of GPU’s spread out across many machines. Another way to get a big accuracy improvement is to decode the CTC probability matrix using a Language Model and the CTC beam search algorithm. CTC type models are very dependent on this decoding process to get good results. Luckily there is a handy that allows you to do that. open source library This tutorial was made to be more accessible so it’s a relatively small model (23 million Parameters) compared to something like BERT (340 million Parameters). It seems to be the larger you can get your network, the better it performs, although there are diminishing returns. A larger model equating to better performance is not always the case though, as proven by OpenAI’s research . Deep Double Descent This model has 3 residual CNN layers and 5 Bidirectional GRU layers which should allow you to train a reasonable batch size on a single GPU with at least 11GB of memory. You can tweak some of the hyper parameters in the main function to reduce or increase the model size for your use case and compute availability. Latest Advancements In Speech Recognition with Deep Learning Deep learning is a fast-moving field. It seems like you can’t go a week without some new technique getting state of the art results. Here are a few of things worth exploring int the world of speech recognition. Transformers Transformers have taken the Natural Language Processing world by storm! First Introduced in the paper , transformers have been taking and modified to beat pretty much all existing NLP task dethroning RNN’s type architectures. The Transformer’s ability to see the full context of sequence data is transferable to speech as well. Attention Is All You Need Unsupervised Pre-training If you follow deep learning closely you’ve probably heard of BERT, GPT, and GPT2. These Transformer models have first pertained on a language modeling task with unlabeled text data, and fine-tuned on a wide array of NLP task and get state of the art results! During pre-training, the model learns something fundamental on the statistics of language and uses that power to excel at other tasks. We believe this technique has great promises on speech data as well. Word Piece Models Our model defined above output characters. Some benefits to that are the model doesn’t have to worry about out of vocabulary words when running inference on speech. So for the word each character has is its own label. The downside to using characters are inefficiency and the model being prone to more errors because you’re predicting one character at a time. c h a t Using the whole word as labels have been explored, to some degree of success. Using this method, the entire word would be the label. But using whole words, you would have to keep an index of all possible vocabularies to make a prediction, which is memory inefficient with the possibility of running into out of vocabulary words during prediction. The sweet spot would be using word piece or sub-word units as labels. chat Instead of characters for the individual label, you can chop up the words into sub-word units, and use those as labels, i.e. . This solves the out of vocabulary issue, and is much more efficient, as it needs fewer steps to decode then using characters, and without the need to have an index of all possible words. Word pieces have been used successfully with many NLP models, like BERT and would work natural with speech recognition problems as well. ch at Previously published at https://www.comet.ml/site/customer-case-study-building-an-end-to-end-speech-recognition-model-in-pytorch-with-assemblyai/