在变量序列批次中采样softmax损失？

Question

背景信息：我正在研究序列到序列模型，现在我的模型接受可变长度输入张量（非列表），其输入形状对应于[批量大小，序列长度]。但是，在我的实现中，序列长度是未指定（设置为None）以允许可变长度输入。具体而言，输入序列批次仅填充到该批次中最长序列的长度。这大大加快了我的训练时间，所以我宁愿保持这种方式，而不是回到划线模型和/或将训练数据中的所有序列填充到相同的长度。我正在使用TensorFlow 1.0.0。

问题：我目前正在使用以下方法来计算损失（运行得很好）。

loss = tf.losses.sparse_softmax_cross_entropy(
    weights=target_labels,  # shape: [batch size, None]
    logits=outputs[:, :-1, :], # shape: [batch size, None, vocab size]
    weights=target_weights[:, :-1]) # shape: [batch size, None]

其中词汇大小通常约为40,000。我想使用一个采样softmax，但我遇到了一个问题，这是由于输入形状的未指定性质。根据{{​​3}}，，它需要为每个时间段单独输入输入。但是，我不能打电话，例如，

tf.unstack(target_labels, axis=1)

因为轴事先是未知的。有谁知道我可能会如何实现这个？可以假设，由于dynamic_rnn和documentation for tf.nn.sampled_softmax_loss似乎都没有这样做的问题，因此可以通过某种方式使用采样的softmax损失来实现变通方法。在挖掘了源代码甚至模型库之后，我空手而归。任何帮助/建议将不胜感激。