是否可以将 CuDNNLSTM 与 Google Colab 的 TPU 一起使用？

【问题标题】：Is it possible to use CuDNNLSTM with Google Colab's TPU?是否可以将 CuDNNLSTM 与 Google Colab 的 TPU 一起使用？
【发布时间】：2020-12-16 21:00:08
【问题描述】：

我可以使用他们的 GPU 来做到这一点，但使用他们的 TPU 时它会检索到一个错误... 请问这里有人知道我想念吗？将 TPU 与 CuDNNLSTM 一起使用是否有意义？还是 CuDNNLSTM 只是为 GPU 量身定制的？提前非常感谢。

【问题讨论】：

TPU 不支持 CUDA 或 CuDNN。它将需要不同的软件堆栈。我正在删除cuda 标签，因为关于如何在 TPU 上运行某些东西的问题实际上与 CUDA 无关。

标签： google-colaboratory tpu

【解决方案1】：

keras.layers.CuDNNLSTM 仅在 GPU 上受支持。但在 Tensorflow 2 中，内置 LSTM 和 GRU 层已更新为在 GPU 可用时默认利用 CuDNN 内核。

以下是来自Performance optimization and CuDNN kernels的详细信息：

在 TensorFlow 2.0 中，内置 LSTM 和 GRU 层已更新为在 GPU 可用时默认利用 CuDNN 内核。通过此更改，之前的 keras.layers.CuDNNLSTM/CuDNNGRU 层已被弃用，您可以构建模型而无需担心它将运行的硬件。

您可以只使用内置的 LSTM 层：tf.keras.layers.LSTM，它可以在 TPU 和 GPU 上运行。

【讨论】：

太棒了！我在 tf 2 中没有意识到这一点感谢您的输入：它很有用。