系统信息

• 是否编写了自定义代码(与在TensorFlow中使用的库存示例脚本相反):是的
• OS平台和发行版(例如，Linux Ubuntu 16.04):Ubuntu 20.04
• 从哪里安装的TensorFlow(源代码或二进制文件):二进制(最新GPU docker)
• TensorFlow版本(请使用以下命令):v2.5.0-rc3-213-ga4dfb8d1a71 2.5.0(但这也可以在较早的版本中重现)。
• Python版本：3.6.9
• CUDA/cuDNN版本：11.2 / 8.1.0
• GPU型号和内存：NVIDIA Quadro P1000
  描述当前行为

如您在此处的代码中所看到的，CUDNN LSTM和GRU的实现方式是通过向图中添加两个具有相同名称的函数。Grappler的ImplementationSelector优化器稍后运行并选择正确的实现。然而，当使用tf.data.experimental.scan时，根据我之前看到的内容(我不完全记得很清楚),该函数会被重命名，这会导致实现选择器失败。
实际上，非CUDNN版本是默认实现，因此会回退到那个实现。

描述预期行为

层应在可能的情况下使用CUDNN运行。

Contributing

• 您是否想提交PR?(是/否):这可能超出了我的能力范围，但如果解决方案很简单，我可以实现它。
• 简要描述您的候选解决方案(如果有贡献的话):
  独立代码以重现问题

不幸的是，我想不出一种不使用Tensorboard就能显示此问题的方法，但这里是：

import tensorflow as tf

model = tf.keras.models.Sequential([
    tf.keras.layers.GRU(4),
])

@tf.function
def foo(dataset):
    a = 1.
    # Autograph converts this to tf.data.experimental.scan
    for i in dataset:
        a = model(i)
    return a # return the value to avoid pruning

inp_tensor = tf.random.normal([32, 10, 8])
dataset = tf.data.Dataset.from_tensor_slices([inp_tensor])

tf.profiler.experimental.start("/tmp/tensorboard")

# This successfully uses CUDNN
model.predict(inp_tensor)

# This doesn't use CUDNN
foo(dataset)

tf.profiler.experimental.stop()

在这里，您可以看到model.predict使用了CUDNN:

而在这里，您可以看到foo(dataset)没有使用CUDNN: