gpu上的tensorflow 2.0自定义图层

Question

TensorFlow中完全定制的图层会自动在GPU上运行吗？我注意到在这个文档（https://www.tensorflow.org/guide/keras/rnn#rnn_layers_and_rnn_cells）中，RNN包装器似乎不会使用CudNN？这意味着它不会在GPU上运行，对吗

Answer 1

您的自定义层仍将使用GPU，您可以确认这一点，如本answer中所述

虽然自定义层不会使用cuDNN，但您是对的。为什么这很重要？在NVidia之后引用：

cuDNN provides highly tuned implementations for standard routines such as forward and backward convolution, pooling, normalization, and activation layers

换句话说，使用这些优化的原语将提高培训的性能。在cuDNN: Efficient Primitives for Deep Learning文件中提供了许多详细解释的例子。以空间卷积为例。未优化的实现将使用“幼稚”的方法，而cuDNN使用各种技巧来减少操作数量并适当地批处理它们。与传统CPU相比，GPU仍然很快，cuDNN只是让它更快。有关更近期的独立基准测试，请查看例如this article

尽管如此，如果Tensorflow在GPU模式下运行，完整的计算图形将在GPU上执行（据我所知，甚至没有简单的方法可以取出图形的一部分，即中间层，然后放到CPU上）