tensorflow-gpu - TensorFlow 如何使用 cudnn

Question

我目前正在阅读 Tensorflow 源代码，并对内核的实现感到好奇。我发现大多数 gpu 实现都指向 Eigen。谁能告诉我 tensorflow 如何通过 Eigen 或其他方式使用 cuDNN？

Answer 1

是的，大多数基本内核都使用使用普通 CUDA 的 Eigen。使用 cuDNN（例如卷积）的内核通过此集成：https ://github.com/tensorflow/tensorflow/tree/master/tensorflow/stream_executor/cuda

这是一个示例 Conv 内核，它检索支持的 Conv 算法（包括 cuDNN，如果它已链接且可用），运行并选择最佳算法，最后使用它。