tensorflow 2.0 中是否有 cudnnLSTM 或 cudNNGRU 替代方案

Question

该CuDNNGRU中TensorFlow 1.0是非常快。但是当我转移到TensorFlow 2.0我无法找到CuDNNGRU. 简单GRU在 TensorFlow 2.0.

有没有办法使用CuDNNGRU的TensorFlow 2.0？

Answer 1

在导入的实现已经过时-相反，LSTM并且GRU将默认为CuDNNLSTM和CuDNNGRU，如果所有条件都满足：

1. activation = 'tanh'
2. recurrent_activation = 'sigmoid'
3. recurrent_dropout = 0
4. unroll = False
5. use_bias = True
6. 输入，如果 masked，严格右填充
7. reset_after = True （仅限 GRU）

还要确保 TensorFlow 使用 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()))

更新：在 Colab 上运行时，TF 2.0.0 似乎存在问题，使 CuDNN 正常工作；试试吧!pip install tensorflow==2.1.0。