使用GPU支持在Tensorflow中创建自定义LSTMCell

Question

我想在我正在创建的LSTM模型中整合注意力成分。不幸的是，在我使用的tensorflow 2.3.1中，如果您将LSTMCell子类化，那么您必须在CPU上运行该模型。从tensorflow文档中：

CuDNN is only available at the layer level, and not at the cell level.

这意味着如果我尝试这样的方法，我将被降级到CPU：

output=keras.layers.RNN(AttentionLSTMCell(400), return_sequences=True, stateful=False)(input_layer);

其中AttentionLSTMCell是我创建的一个自定义类，它将接受一些附加常量（通常是前一个时间戳的输出和一些新的输入），这些常量将调节LSTM的输出。事实上，文件似乎表明，即使只允许特定条件。我将深入研究创建一个完整的自定义层（可能复制现有的层，看看是否可以在调用中添加新的输入），但是有更好的方法吗？这使得原型制作相当困难。大型递归网络的训练速度很慢，特别是在我集成图像数据作为输入的情况下