Search code examples
pythontensorflowkerasrecurrent-neural-networktensorflow2.0

Is there cudnnLSTM or cudNNGRU alternative in tensorflow 2.0

The CuDNNGRU in TensorFlow 1.0 is really fast. But when I shifted to TensorFlow 2.0 i am unable to find CuDNNGRU. Simple GRU is really slow in TensorFlow 2.0.

Is there any way to use CuDNNGRU in TensorFlow 2.0?

Solution

  • The importable implementations have been deprecated - instead, LSTM and GRU will default to CuDNNLSTM and CuDNNGRU if all conditions are met:

    1. activation = 'tanh'
    2. recurrent_activation = 'sigmoid'
    3. recurrent_dropout = 0
    4. unroll = False
    5. use_bias = True
    6. Inputs, if masked, are strictly right-padded
    7. reset_after = True (GRU only)

    Also ensure TensorFlow uses GPU:

    import tensorflow as tf
from tensorflow.python.client import device_lib
print(device_lib.list_local_devices())
print('Default GPU Device: {}'.format(tf.test.gpu_device_name()))

    Update: there appears to be a problem w/ TF 2.0.0 when running on Colab in getting CuDNN to work; try !pip install tensorflow==2.1.0 instead.