Tensorflow.js tf.layers.rnn() Function

Tensorflow.js is an open-source library developed by Google for running machine learning models and deep learning neural networks in the browser or node environment. 

Tensorflow.js tf.layers.rnn() function is basically base class for recurrent layers. 

Syntax:

tf.layers.rnn( args );

Parameters: 

• args: It is an object with the following fields:
  • cell: It should be an instance of an RNN cell or an array of instances of RNN cells.
  • returnSequences:  It should be boolean. It defines which both of these weather last output or full sequence return in the output sequence.
  • returnState: It should be boolean. It defines the return last state in output. 
  • goBackwards: It should be boolean. It defines whether to process the input in reverse and reverse sequence or not. 
  • stateful: It should be boolean. It defines whether to use the last state of the batch as the initial state for current batch or not. 
  • unroll: It should be boolean. It tells whether to use unrolled network else, use a symbolic loop. 
  • inputDim: It should be a number. This is used when this layer is used as the first layer of the model. It defines the dimension of input in the layer. 
  • inputLength: It should be a number. This field should be when the sequence in input data is constant. 
  • inputShape: It should be an array of numbers. This field is used to create input layer which is used to be inserted before this layer. 
  • batchInputShape: It should be an array of numbers. This field will be used if inputShape and this field are provided as a parameter for creating the input layer which is used to insert before this layer. 
  • batchSize: It should be a number. In the absence of batchInputShape this field is used to create batchInputShape with inputShape. batchInputShape : [ batchSize ,  …inputShape].
  • dtype: If this layer is used as the input layer, then this field is used as the data type for this layer. 
  • name: It should be a string type. this field defines the name for this layer. 
  • trainable: It should be boolean. This field defines whether the weights of this layer are trainable with fit or not.
  • weights: This should be a tensor that defines the initial weight value for this layer.
  • inputDType: This is the data- type which is used for Legacy support.

Return Value: It returns RNN.

This layer supports masking for input data with any number of timesteps. We can set RNN layers to be ‘stateful’, which means that the state computed for the samples in one batch will be reused as initial states for the next batch. 

Input shape should be a 3D tensor with shape [ batchSize, timeSteps, inputDim ]

As per the parameters and values provide to this layer we get the shape of the output layer: 

Example 1: If returnState is set true then this function return array of tensor in which the first tensor is the output. the remaining tensors are the state at the last time steps. The Shape of all the tensors in the array will be [ batchSize, units].

Javascript

// Cells for RNN
const cells = [
   tf.layers.lstmCell({units: 4}),
   tf.layers.lstmCell({units: 8}),
];
const rnn = tf.layers.rnn({cell: cells, returnState:true });

// Create an input with 8 time steps and
// 16 length vector at each step.
const Input_layer = tf.input({shape: [8, 16]});
const Output_layer = rnn.apply(Input_layer);

console.log(JSON.stringify(Output_layer[0].shape));

Output:

​[null, 8]

Example 2: If returnState is not set and returnSequence is set to true then the Output tensor will have the shape: [ batchSize, timeSteps, units ].

Javascript

// Cells for RNN
const cells = [
    tf.layers.lstmCell({units: 16}),
    tf.layers.lstmCell({units: 32}),
];
const rnn = tf.layers.rnn({
    cell: cells,
    returnState: false,
    returnSequences: true
});

// Create an input with 8 time steps and
// 16 length vector at each step.
const Input_layer = tf.input({shape: [4, 8]});
const Output_layer = rnn.apply(Input_layer);
console.log(JSON.stringify(Output_layer.shape));

Output:

[null,4,32]

Example 3: If both returnState and returnSequences are not defined then the shape of output will be [ batchSize, units ].

Javascript

// Cells for RNN
const cells = [
   tf.layers.lstmCell({units: 4}),
   tf.layers.lstmCell({units: 8}),
];
const rnn = tf.layers.rnn({cell: cells});

// Create an input with 10 time steps and
// 20 length vector at each step.
const input = tf.input({shape: [10, 20]});
const output = rnn.apply(input);

console.log(JSON.stringify(output.shape));

Output:

[null,8]

References: https://js.tensorflow.org/api/latest/#layers.rnn