Python Tensorflow图像字幕/seq2seq解码器作为模型

我使用一个序列预测模型，从编码输入的潜在表示开始，它形成解码器的初始状态。它可以是来自图像（用于字幕）的特征向量，也可以是seq2seq编码器的结果

我的模型是在老师的强制下训练的，而且速度很快。然而，推理速度非常慢，因为我以（伪代码）的形式进行了逐步序列扩展

我在这个阶段做了很多优化，所以我可以并行预测数百个查询的序列，但1）在CPU上运行时，它的线性扩展为>32个左右的序列，2）在GPU上运行时实际上比在CPU上运行时慢，这可能是因为数据在每一步之后都必须来回移动，GPU速度没有任何好处

另外，我还使用了一个非贪婪序列搜索，它不是Beam搜索，但可以以类似于a*的方式回溯，大致如下（伪代码）：

这使得top-n序列从最后预测的位置回溯。同样，这在并行预测的CPU方面或多或少是优化的

我认为为了更快，我需要在GPU上完全运行预测，而无需将数据移回。但我不知道如何用TensorFlow写这个。有一个

tfa.seq2seq

（以前的

tf.contrib.seq2seq

），它有一个解码器的基础结构，可以作为模型高效运行，但我找不到太多文档

请注意，我的模型（Keras functional API；它也可以与

model（）

一起使用，而不是

model.predict（）

或者我可以将输出张量连接到其他地方）不仅仅是3个LSTM层，而是具有一些有状态的内联特性工程，因此需要在模型中完成

tfa.seq2seq.Decoder

似乎期望一个单元将其自身包裹起来

问题: 1） 对于独立于

tfa.seq2seq

体系结构构建和训练的黑盒模型，我可以使用

tfa.seq2seq

解码器体系结构吗？如果是，我在哪里可以找到相关信息？ 2） 有没有关于如何直接在tensorflow上实现贪婪和非贪婪序列搜索的指南，而不必回到上面的python代码？我知道我可能不得不放弃我的非贪婪方法，只使用beam搜索，它的性能可能也差不多

sequence_terminated = False
sequence = np.array((0, output_features))
while not sequence_terminated:
   seq_output, seq_states = model.predict(seq_input)
   next_input, sequence_terminated = f(seq_output)
   sequence = np.concatenate(sequence, seq_output)

sequence_terminated = False
sequence = np.array((0, output_features))
states = np.array((0, state_size))
pos = []
from_position = 0
remaining_sequences = 5
while remaining_sequences > 0:
   seq_output, seq_states = model.predict(seq_input)
   sequence = np.concatenate(sequence, seq_output)
   states = np.concatenate(states, seq_states)
   pos.append(from_position)
   # Based on the outputs until now, find what sequence stub to continue:
   next_input, next_states, from_position, terminated = f(sequence, states)
   if terminated:
      remaining_sequences = remaining_sequences - 1