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

Question

我可以使用他们的 GPU 执行此操作，但使用他们的 TPU 时会检索到一个错误...请问这里有没有人知道我想念什么？将 TPU 与 CuDNNLSTM 实际结合使用是否有意义？还是 CuDNNLSTM 只是为 GPU 量身定做的？非常感谢。

Answer 1

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

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

In TensorFlow 2.0, the built-in LSTM and GRU layers have been updated to leverage CuDNN kernels by default when a GPU is available.With this change, the prior keras.layers.CuDNNLSTM/CuDNNGRU layers have been deprecated, and you can build your model without worrying about the hardware it will run on.

您可以只使用内置的 LSTM 层:tf.keras.layers.LSTM它适用于 TPU 和 GPU。