Tacotron 2

October 25, 2020

---

Tacotron 2 is a Text-to-speech (TTS) model developed by Google. It consists of two components, a recurrent sequence-to-sequence architecture with attention that generates a magnitude mel-spectrogram from a sequence of characters, and a modified WaveNet which generates audio (time-domain waveform) from the mel spectrogram. The second component is called a vocoder. In this post, we will only look at the training of the first network to generate the mel spectrograms. We will use the Griffin-Lim implementation in librosa to obtain the audio from the mel spectrograms.

A single training iteration (forward and backward propagation) invokes about 12,000 GPU CUDA kernels for every second of audio! Fortunately, there is a lot of repetition, due to the architecture being recurrent. The best way to understand a DL network and GPU performance is to understand every single CUDA kernel i.e. which layer of the network invoked the kernel, with what arguments (tensor shapes and datatypes) and in which direction (forward or backward).

In this post, we will categorize every kernel used in the training of Tacotron 2. All the information in the tables below was obtaining using Nvidia's PyTorch Profiler, PyProf, on a Turing T4 GPU. The information below is only a subset of what is provided by PyProf. The code for Tacotron 2 was obtained from here. The annotated code and instructions for obtaining a detailed profile are here. Note that different GPUs will have slightly different kernel names e.g. volta_* as opposed to turing_*.

Model Architecture

For profiling, we used a single short sentence "In his defense" (LJ011-0028 from LJ Speech Dataset) repeated 64 (batch_size) times. The length of the corresponding audio is about 1 second. The table below shows the output mel spectrogram and the corresponding audio (obtained using Griffin-Lim) at various points during the training. The parameters in the code are here.

Model Parameters

batch_size = 64

# Character Embedding
n_symbols = 148
symbols_embedding_dim = 512

# Encoder Convolution
encoder_kernel_size = 5
encoder_n_convolutions = 3
encoder_embedding_dim = 512

# Audio Parameters
sampling_rate = 22050
filter_length = 1024
hop_length = 256
win_length = 1024
n_mel_channels = 80
mel_fmin = 0.0
mel_fmax = 8000.0

# Decoder PostNet
postnet_embedding_dim = 512
postnet_kernel_size = 5
postnet_n_convolutions = 5

# Location Layer parameters
attention_location_n_filters = 32
attention_location_kernel_size = 31

GPU Kernels

The tables below show the GPU kernels invoked in 1 training step. For every GPU kernel we show the direction (fprop, bprop), name of the layer, name of the operation, and the input tensor shapes / matrix dimensions for the operation. PyProf provides a lot of additional information for every GPU kernel e.g. grid dimensions, block dimensions, silicon time, datatypes, flops, bytes, tensor core usage and so on.

Zero out the gradients.

At the beginning of an iteration, we zero out the gradients.

Kernels 61 and 62 calculate the maximum sentence length in a batch. Kernel 63 copies the target padded mel spectrogram to the GPU which is used to calculate the loss.

Encoder

The network is composed of an encoder and a decoder with attention. The encoder converts a character sequence into a hidden feature representation, which the decoder consumes to predict a spectrogram. The encoder consists of a character embedding layer, 3 1-D convolution layers and a bi-directional LSTM layer.

Encoder: Character Embedding, 3 Convolution Layers

Kernel 64 corresponds to the character embedding. Note that the input sentence "In his defense" has 14 characters and the batch_size is 64. The embedding table has the shape [n_symbols, symbols_embedding_dim]. Kernels 65 through 91 correspond to the three convolution layers. For each convolution the input shape [N,C,H,W] is [64,512,1,14], the number of filters K is encoder_embedding_dim i.e. 512, the kernel size [R,S] is [1,encoder_kernel_size] i.e. [1,5] and the padding [ph,pw] is [0,2]. The output shape [N,K,P,Q] is [64,512,1,14]. Each convolution is followed by batch normalization, ReLU and dropout. Kernels 71, 80 and 89 correspond to scalar additions 1+1, and I am not sure of their purpose.

Encoder: Bidirectional LSTM

Kernels 95 through 169 correspond to the bi-directional LSTM layer. It contains 512 units, 256 in each direction. The matrix multiplication (GEMM) in a LSTM cell can be broken down into two components, a recurrent component which depends on the hidden state H, and a non-recurrent component which depends on the input X. The recurrent component has to be performed sequentially while the non-recurrent component can be done in parallel. In addition, the non-recurrent component for both directions can be combined. In our input, the sequence length (number of characters in the input sentence) is 14, which results in 14 recurrent GEMMs in each direction plus 14 non-recurrent GEMMs for a total of 42 GEMM kernels. These kernels have the name volta_fp16_* in the table below. A LSTM cell has many pointwise (elementwise) operations e.g. sigmoid, tanh etc. With our input, we get 28 LSTM_elementWise_fp kernels, 14 in each direction. Kernels 93, 94 and 170 correspond to pack_padded_sequence, flatten_parameters, and pad_packed_sequence function calls respectively. These functions are not strictly required but using them with a GPU can result in faster execution.

Decoder

The decoder is an autoregressive (previous output frame is used to compute the current output frame) recurrent neural network which predicts a mel spectrogram, one frame at a time. This autoregressive procedure is followed during inference. However, during training, we feed the target previous frame (available from labelled data), instead of the frame generated by the network, to compute the current output frame.

Decoder: Prepare inputs (from labelled data)

Kernel 171 corresponds to getting an all zeros frame to use as the first decoder input. The shape of a frame is [batch_size,n_mel_channels] i.e. [64,80]. Kernel 172 corresponds to concatenating the first all zeros frame to the labelled data (target). The shape of the labelled data is [frames,batch_size,n_mel_channels] i.e. [98,64,80]. The number of frames depends on the length of the audio (in seconds), sampling_rate, hop_length and win_length.

Decoder: PreNet

During inference, the PreNet is the first stage of the decoder and is called in an autoregressive fashion i.e. for every frame. During training, instead of calling the PreNet for every labelled (target) frame, a performance optimization is to call the PreNet just once, for all the frames generated above. The PreNet is a simple two layer MLP, each with 256 units, ReLU activation and dropout. The output of the PreNet (one frame at a time) is then used to generate the target frames (one frame at a time).

Decoder: State initialization

Kernels 179 through 183 calculate the attention mask. Kernels 184 through 190 zero out the hidden state of the first LSTM, the cell state of the first LSTM, the hidden state of the second LSTM, the cell state of the second LSTM, the initial attention weights, the cummulative attention weights and the attention context respectively.

Tacotron uses location sensitive attention. The equation for the scoring mechanism (Eq. 9 in the paper) is given by,

Here, $i$ refers to the $i$-th frame (decoder step) and $j$ is an index into the encoder state. The second term of the equation, $Vh_j$, does not depend on the decoder step and can be computed ahead of time for all $j$. The encoder state, $h$, has the shape [batch_size,sequence_length,lstm_hidden_state] i.e. [64,14,512]. Kernels 191 and 192 correspond to the operation, $Vh$.

Decoder: LSTM, Location Sensitive Attention, Linear Projection, Stop Token

Kernels 193 through 226 correspond to predicting one frame of the mel spectrogram. Since, our example input has 98 frames, we get a total of $34 \times 98 = 3332$ kernels. The output of the PreNet (one frame) is concatenated with the attention context vector. Kernel 193 corresponds to this operation. The concatenated tensor is fed to a LSTM cell with 1024 hidden units with dropout. Kernels 194-197 correspond to this operation. Kernels 198-213 correspond to the location sensitive attention.

The scoring function proposed in content based attention, is given by,

Location sensitive attention extends this by taking into account the alignment (a.k.a attention weights, $\alpha$) produced by the previous step and the cummulative attention weights (running sum of the previous attention weights). The attention weights and cummulative attention weights have the shape [batch_size,sequence_length] i.e. [64,14]. The attention and cummulative attention weights are unsqueezed and concatenated in kernel 198.

Location features are computed using 32 1D convolution filters of length 31 (kernels 201-204). The input shape [N,C,W] is [batch_size,2,sequence_length] i.e. [64,2,14], the number of filters K is attention_location_n_filters i.e. 32, the filter size [R,S] is [1,attention_location_kernel_size] i.e. 31, the padding [ph,pw] is [0,31] so that Q = W = sequence_length. The output shape [N,K,Q] is [64,32,14] and corresponds to $f_i$ in the equation below.

The scoring function used in location sensitive attention is given by,

Here, $i$ refers to the $i$-th frame (decoder step) and $j$ is an index into the encoder state. The third term of the equation, $Uf_i$, is computed in kernels 205 and 206. The first term of the equation, $Ws_{i-1}$, uses the output of the LSTM cell and is computed in kernels 199 and 200. Recall, that the second term of the equation was computed earlier.

Once, we have all the three terms, we add them up (kernels 207 and 208). Kernels 209 and 210 correspond to tanh and matrix-vector multiplication. This gives us the attention scores. Since sentences in a batch have different lengths, kernels 211 applies a mask to the attention scores. Kernel 212 corresponds to the softmax operation which normalizes the scores. Using the terminology from the paper, we calculate the glimpse, $g_i$, in kernel 213, which is the attention context vector.

In kernel 214, we accumulate the new attention weights which will be used in the next step (frame). In kernel 215, we concatenate the output of the first LSTM and the attention context and feed it to another LSTM cell with 1024 hidden units with dropout (kernels 216-219).

Finally, we concatenate the LSTM output and the attention context (kernel 220). This output is passed through a linear layer (kernels 221-223) to produce a spectrogram frame, whose shape is [batch_size,n_mel_channels] i.e. [64,80]. In parallel, the concatenated output is passed through another linear layer (kernels 224-226) to predict the probability that the output sequence has completed i.e. a stop token. The output of this linear layer has the shape [batch_size,1] i.e. [64,1].

Decoder: Output Stacking

At every step, the decoder outputs a mel-spectrogram frame of shape [batch_size,n_mel_channels] i.e. [64,80] and a stop token of shape [batch_size] i.e. [64]. In kernels 3526 and 3527, we stack the stop tokens from all steps and make them contiguous to give us a tensor of shape [batch_size,frames] i.e. [64,98]. Likewise, in kernels 3528 and 3529, we stack the spectrogram frame from all steps and make them contiguous to give us a tensor of shape [batch_size,frames,n_mel_channels] i.e. [64,98,80].

During every step, we calculate the alignment (attention weights) which has the shape [batch_size,sequence_length]. In kernel 3525, we stack the alignments from all the steps. This operation is not required for training and is probably used for visualization.

Decoder: PostNet

The stacked and contiguous mel spectrograms from the previous step are fed to the PostNet. The PostNet block has 5 convolution layers with a residual connection. Each convolution is followed by batch normalization, tanh activation and dropout. Tanh activation is present on all but the last layer. Kernels 3530 through 3575 correspond to the five convolution layers. The input to the first convolution [N,C,H,W] is [batch_size,n_mel_channels,1,frames] i.e. [64,80,1,98], the number of filters K is 512, the kernel size [R,S] is [1,5] and the padding [ph,pw] is [0,2]. The output [N,K,P,Q] is [64,512,1,98]. For the second, third and fourth convolution both the input and output shape is [64,512,1,98]. The number of filters, kernel size, padding is the same as above. For the fifth convolution, the output shape is [64,80,1,98]. The number of filters is 80. The kernel size and padding is the same as above. Kernel 3576 corresponds to the residual connection.

Output Postprocessing

Similar to kernels 179-183, kernels 3577-3581 calculate the mask. Kernels 3582, 3583, 3584 apply the mask to the mel spectrogram before the PostNet, the mel spectrogram after the PostNet and to the stop tokens respectively. Kernels 3585-3588 convert the datatype of the mel spectrogram before the PostNet, mel spectrogram after the PostNet, stop token and alignment tensors respectively from float16 to float32 before feeding it to the loss layers. The alignment tensor is not used for loss calculation and therefore, kernel 3588 is not required.

Loss

Kernels 3589 and 3590 calculate the MSE loss due to the mel spectrogram before the PostNet. Kernels 3591 and 3592 calculate the MSE loss due to the mel spectrogram after the PostNet. Kernel 3593 sums up the two losses. Kernels 3594-3608 correspond to the BCE loss from the stop token prediction. Kernel 3609 sums up all the losses. Kernel 3610 most likely corresponds to loss scaling by a factor.

Back propagation: Loss

Kernels 3612 through 3620 correspond to back propagation through the loss layer.

Kernels 3621-3623 most likely correspond to back propagation through the convert operations (kernels 3585-3587).

Back propagation: PostNet

Kernels 3624 through 3672 correspond to back propagation through the five convolution layers in PostNet.

Back propagation: LSTM, Location Sensitive Attention, Linear Projection, Stop Token

Our example input has 98 frames and therefore we perfrom back propagation through the 2 LSTM layers with location sensitive attention and linear projection layers, 98 times. The last frame invokes 53 kernels, the first frame invokes 66 kernels and all the intermediate frames invoke 78 kernels for a total of $53 + 96\times 78 + 66 = 7607$ kernels. I am not sure why the last and first frame behave differently. The table below shows the kernels invoked during back propagation for one of the intermediate frames with 78 kernels. For many kernels, we cannot ascertain any information except the name and are shown as na.

Back propagation: PreNet

Kernels 11284 through 11293 correspond to back propagation through the PreNet block.

Back propagation: Bi-directional LSTM

Kernels 11294 through 11402 correspond to back propagation through the bi-directional LSTM layer. During back propagation, we have to calculate $\partial H$, $\partial X$ (data gradients) and $\partial W$ (weight gradient). The number of data gradients to calculate is equal to sequence length * 2 because we have a bi-directional LSTM. The 28 triplets of LSTM_elementWise_bp1, cutlass::Kernel, splitKreduce_kernel most likely correspond to this operation. It appears contributions to $\partial W$ are calculated every 4 steps in each direction. The number of (batched) GEMM kernels required to calculate the weight gradients is equal to $\left \lceil \frac{\text{sequence length}}{4} \right \rceil \times 2 = 8. $The kernels turing_*_nn most likely correspond to this operation. There are 6 additional GEMM kernels (volta|turing)_*_nt, which I am unable to decipher.

Back propagation: 3 Convolution Layers

Kernels 11403 through 11432 correspond to back propagation through the 3 convolution layers.

Back propagation: Character Embedding

Kernels 11433 through 11444 correspond to back propagation through the embedding layer.

Loss scaling

Kernels 11445-11448 most likely correspond to reverting the loss scaling of the gradients.

Clip by L2 norm

The model has 60 parameters and kernels 11449 through 11508 correspnd to calculating the L2 norm. We show only 1 entry in the table below. In kernels 11509 and 11510, we stack the L2 norms and calculate the norm again. Kernels 11511-11514 do some checks and calculate the scaling factor. Kernels 11515 through 11574 scale the gradients. We show only 1 entry in the table below.

Adam optimizer

The last step is to apply the Adam optimizer on the model weights. The model has 60 parameters. Each call to the Adam optimizer invokes 9 kernels, for a total of 540 kernels (11575 through 12114). This is not an optimized implementation and one can use the fused Adam implementation from Nvidia Apex. The table below shows the kernels invoked for just 1 parameter update, the character embedding table.

Scaling

Kernels 12115 and 12116 appear to be doing some sort of scaling but am not sure why.