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

Question

TensorFlow 1.0里的CuDNNGRU真快。但是当我转到 TensorFlow 2.0 时，我找不到 CuDNNGRU。简单 GRU 在 TensorFlow 2.0 中真的很慢。

有什么方法可以在TensorFlow 2.0中使用CuDNNGRU吗？

Answer 1

可导入的实现已被弃用 - 相反，LSTM 和 GRU 将默认为 CuDNNLSTM 和 CuDNNGRU 如果所有 conditions are met:

1. activation = 'tanh'
2. recurrent_activation = 'sigmoid'
3. recurrent_dropout = 0
4. unroll = False
5. use_bias = True
6. 输入，如果被屏蔽，则严格右填充
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 上让 CuDNN 工作时，TF 2.0.0 似乎存在问题；请尝试 !pip install tensorflow==2.1.0。