tensorflow - 你怎么知道你在 TensorFlow 2.0 中使用的是 LSTM 或 GRU 的 CuDNN 版本？

Question

TensorFlow 2.0 的一个很好的特性：新的 LSTM 层现在自动决定是否使用 CuDNN 实现（GRU层也是如此）。但是它们的使用不仅仅依赖于 GPU；还有其他需要设置的参数（例如，unroll=False）。

我正在使用带有 GPU 的 tensorflow-gpu（不是一个非常强大的 GPU），而且我的 RNN 像糖蜜一样慢。我想确定我正在运行 CuDNN 实现。日志中有什么地方可以看吗？

我还使用了 tf.compat.v1.keras.layers.CuDNNLSTM，这对我来说同样慢。我可以确定这是使用 CuDNN 实现，还是在某些条件下它可能会退回到其他东西？

标签： tensorflowkeraslstmrecurrent-neural-network

Answer 1

如果您使用的是tf.compat.v1.keras.layers.CuDNNLSTM，那肯定是在使用 cuDNN 实现。里面没有退路。