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

我想在我正在创建的LSTM模型中集成一个注意力成分。不幸的是,在我使用的TensorFlow2.3.1中,如果您将LSTMCell子类化,则必须在CPU上运行模型。来自tensorflow文档:

这意味着如果我尝试这样的方法,我会被降级到CPU:

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

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