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

Question

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

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

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)

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

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

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

这使得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搜索，它的性能可能也差不多