Parameters In Tensorflow Keras RNN and CUDNN RNN

Introduction

Recurrent Neural Network (RNN) is widely used in AI applications of handwriting recognition, speech recognition, etc. It essentially consists of a series of matrix-vector multiplications and there are two popular gating mechanisms: GRU and LSTM. Tensorflow Keras provides high-level APIs for such operations, which are simple to use and productive, because it can handle the parameter creation, initialization, and other preprocessing before calling the actual libraries, where Tensorflow Keras adopts CUDNN as the backend for GPUs. In comparison, there are demands that people would like to use CUDNN directly in their projects for more efficiency and better control of the data flow. This might need porting the Keras code or simply viewing its outputs as reference. However, Keras and CUDNN takes different means to deal with the parameters, leading to the different layouts of parameters. This often causes a bit of confusion when developers work on both APIs.

This post will check the GRU and LSTM layers from Keras, especially focusing on how the parameters are organized in Keras and what transformations are needed to make the parameters compatible for CUDNN. This post assumes people have sufficient background of RNN. The equations used below are borrowed from the NVIDIA CUDNN documentation.

GRU

In these equations, the parameters consisit of three kernel weights/biases (W_i, W_r, W_h, b_Wi, b_Wr, b_Wh) and three recurrent weights/biases (R_i, R_r, R_h, b_Ri, b_Rr, b_Rh). These two sets of weights/biases pairs indicate that we apply bias addition whether the weight-input multiplication is for the layer input x_t or for the recurrent input h_t-1. In fact, the above equations represent a “double bias” mode categorized by CUDNN and apparently there are two other bias modes that apply bias addition only to layer input or only to recurrent input.

Here I would like to point out the x_t and h_t-1 are both column vectors in the shape of (inputSize, 1) and (hiddenSize, 1) respectively to make the matrix-vector multiplication possible. Therefore, the shape of each W weights is (hiddenSize, inputSize) and the shape of each R weights is (hiddenSize, hiddenSize), while the biases are always in (hiddenSize, 1).

The above discussion is based on CUDNN formula; by constrast, Keras takes a slightly different way to interpret the weight-input multiplication, which is more like x_t^TW^T or h_t-1^TR^T. That way, the weights in Keras are stored in a transposed style compared to CUDNN: the W weights are (inputSize, hiddenSize) and the R weights are (hiddenSize, hiddenSize). As for the specific implementation under the hood, CUDNN stores all weights and biases all together as a single flatten array, while Keras uses three arrays:

1. Kernel weights: a (inputSize, 3 x hiddenSize) matrix containing the concatenated three W weights.
2. Recurrent weights: a (hiddenSize, 3 x hiddenSize) matrix containing the concatenated three R weights.
3. Biases: a (2, 3 x hiddenSize) matrix where one row is the concatenated W biases and another is for R biases.

The following python code is to output the parameters stored in Keras GRU layer:

import tensorflow as tf
import numpy as np
from tensorflow.keras import layers
from tensorflow.python.keras.layers import recurrent_v2

batch_size = 1
seq_len = 2
input_size = 2
hidden_size = 3

tf.random.set_seed(seed=1)
x = tf.random.uniform((seq_len, batch_size, input_size))
gru = layers.GRU(hidden_size, time_major=True,
                 return_sequences=True,
                 kernel_initializer='glorot_uniform',
                 recurrent_initializer='orthogonal',
                 bias_initializer='random_uniform')
y = gru(x)

np.set_printoptions(
    formatter={'float': lambda x: "{0:0.6f}".format(x)})

print("Keras Kernel Weights:", gru.get_weights()[0])
print("Keras Recurrent Weights:", gru.get_weights()[1])
print("Keras Biases:", gru.get_weights()[2])

Here I don’t copy/paste the outputs but visualize them with three colors to represent different types of weights/biases:

• Green for r weights/biases
• Red for i weights/biases
• Yellow for h weights/biases

Keras Kernel Weights:

Keras Recurrent Weights:

Keras Biases:

Moreover, I add some printfs in the backend to trace the inputs/outputs of CUDNN calls. The parameters directly taken in by CUDNN are as below and I also tint them with the three colors.

CUDNN Parameters (Weights and Biases):

By comparing the parameters in Keras and CUDNN, we can observe three differences:

1. CUDNN uses a single flat array, while Keras uses three separate arrays.
2. The weights are transposed between CUDNN and Keras.
3. The i and r weights/biases are swapped between CUDNN and Keras.

So, if one is given some weights/biases extracted from Keras and their task is to use them in CUDNN, they have to perform a series of preprocessing on them, such as slicing, permuting, tranposing, concatenation, etc. Fortunately, we can find a “secret” tool from Tensorflow Keras that can partially help over the preprocessing:

params = recurrent_v2._canonical_to_params(
    weights=[
        gru.get_weights()[0][:, hidden_size:hidden_size * 2],
        gru.get_weights()[0][:, :hidden_size],
        gru.get_weights()[0][:, hidden_size * 2:],
        gru.get_weights()[1][:, hidden_size:hidden_size * 2],
        gru.get_weights()[1][:, :hidden_size],
        gru.get_weights()[1][:, hidden_size * 2:],
    ],
    biases=[
        gru.get_weights()[2][0][hidden_size:hidden_size * 2],
        gru.get_weights()[2][0][:hidden_size],
        gru.get_weights()[2][0][hidden_size * 2:hidden_size * 3],
        gru.get_weights()[2][1][hidden_size:hidden_size * 2],
        gru.get_weights()[2][1][:hidden_size],
        gru.get_weights()[2][1][hidden_size * 2:hidden_size * 3],
    ],
    shape=tf.constant([-1]),
    transpose_weights=True)
print("CUDNN-equivalent Params:", params)

Note, as we can see, though it is a single function call, we still need to do most of the slicing operations and don’t forget to swap the i and r (i.e., the 0th and 1st) weights/biases.

LSTM

It would be a lot easier to understand the LSTM params format when we have learned it in the GRU. The LSTM equations show that there are four kernel weights/biases (W_i, W_f, W_o, W_c, b_Wi, b_Wf, b_Wo, b_Wc) and four recurrent weights/biases (R_i, R_f, R_o, R_c, b_Ri, b_Rf, b_Ro, b_Rc). The shapes of the weights and biases are same with those in GRU case.

Still, the two sets of weights/biases in the equations demonstrate the “double bias” mode used by CUDNN. In contrast, Keras adopts a mode more close to the “single bias” mode that only performs the bias addition once for each equation above. So, considering the transposed weights and “single bias” mode, the three Keras LSTM arrays are:

1. Kernel weights: a (inputSize, 4 x hiddenSize) matrix containing the concatenated four W weights.
2. Recurrent weights: a (hiddenSize, 4 x hiddenSize) matrix containing the concatenated four R weights.
3. Biases: a (4 x hiddenSize) vector containing the concatenated biases.

The following python code is to output the parameters stored in Keras LSTM layer:

lstm = layers.LSTM(hidden_size, time_major=True,
                   return_sequences=True,
                   kernel_initializer='glorot_uniform',
                   recurrent_initializer='orthogonal',
                   bias_initializer='random_uniform')
y = lstm(x)
print("Keras Kernel Weights:", lstm.get_weights()[0])
print("Keras Recurrent Weights:", lstm.get_weights()[1])
print("Keras Biases:", lstm.get_weights()[2])

Here I visualize the outputs with four colors to represent different types of weights/biases:

• Green for i weights/biases
• Red for f weights/biases
• Yellow for o weights/biases
• Blue for c weights/biases

Keras Kernel Weights:

Keras Recurrent Weights:

Keras Biases:

Similarly, I tint the parameters taken in by CUDNN with the four colors.

CUDNN Parameters (Weights and Biases):

By comparing the parameters in Keras and CUDNN, we can still observe three major differences:

1. CUDNN uses a single flat array, while Keras uses four separate arrays.
2. The weights are transposed between CUDNN and Keras.
3. The order of weights/biases are same between CUDNN and Keras, but the bias part is padded with (4 x hiddenSize) zeros to fill the gaps of the “double bias” (CUDNN) and “single bias” (Keras) modes.

Now, we take advantage of the “secret” function to convert the Keras parameters to be CUDNN compatible:

params = recurrent_v2._canonical_to_params(
    weights=[
        lstm.get_weights()[0][:, :hidden_size],
        lstm.get_weights()[0][:, hidden_size:hidden_size * 2],
        lstm.get_weights()[0][:, hidden_size * 2:hidden_size * 3],
        lstm.get_weights()[0][:, hidden_size * 3:],
        lstm.get_weights()[1][:, :hidden_size],
        lstm.get_weights()[1][:, hidden_size:hidden_size * 2],
        lstm.get_weights()[1][:, hidden_size * 2:hidden_size * 3],
        lstm.get_weights()[1][:, hidden_size * 3:],
    ],
    biases=[
        tf.zeros((hidden_size,)),
        tf.zeros((hidden_size,)),
        tf.zeros((hidden_size,)),
        tf.zeros((hidden_size,)),
        lstm.get_weights()[2][:hidden_size],
        lstm.get_weights()[2][hidden_size:hidden_size * 2],
        lstm.get_weights()[2][hidden_size * 2:hidden_size * 3],
        lstm.get_weights()[2][hidden_size * 3:hidden_size * 4],
    ],
    shape=tf.constant([-1]),
    transpose_weights=True)
print("CUDNN-equivalent Params:", params)

Note, for LSTM, we don’t need to change orders of the four weights/biases but do need to pad one row of biases to be zeros.

Reference

• NVIDIA CUDNN Documentation: cudnnRNNMode_t
• Tensorflow Keras Layers GRU
• Tensorflow Keras Layers LSTM

• 2022 3
• 2021 8
• 2020 1