I have some TensorFlow 1 code which implements a GRU layer, and I am updating it to TensorFlow 2. So instead of the hand-written layer iterating over timesteps with a GRUCell I am using the in-built TensorFlow GRU layer. However I could not find a way to specify multiple layers using that class. The PyTorch GRU layer has a construction parameter
num_layers which allows for stacking multiple GRUs:
dim = 1024 gru = torch.nn.GRU( input_size=dim, hidden_size=dim, num_layers=2 )
TensorFlow 1 has the MultiRNNCell, which enables one to construct a 'stack' of rnn cells which then behave as a single cell:
num_rnn = 4 dim = 1024 gru_stack = tf.contrib.rnn.MultiRNNCell( [tf.contrib.rnn.GRUCell(dim) for _ in range(num_rnn)])
In TensorFlow 2 I could use the StackedRNNCells class, but that's just a replacement for MultiRNNCell and does the same thing. And it still defines only a single timestep for an RNN layer.
I want to use the GRU layer because that has in-built support for CuDNN and better performance (and it requires less code). It seems the only option is to manually define multiple GRUs in a sequence. Am I right? I guess this is not a big deal, but it's slightly awkward if (as in my case) you wish to expose the number of GRU layers as a parameter for a custom layer.