python - 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如果全部conditions are met :

activation = 'tanh'

recurrent_activation = 'sigmoid'

recurrent_dropout = 0

unroll = False

use_bias = True

输入，如果被屏蔽，严格右填充

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反而。