machine-learning - 可变序列批次的采样softmax损失？

Question

背景信息：我正在研究序列到序列模型，现在我的模型接受可变长度输入张量（不是列表），其输入形状对应于 [批量大小，序列长度]。但是，在我的实现中，未指定序列长度（设置为无）以允许可变长度输入。具体来说，输入序列批次仅填充到该批次中最长序列的长度。这大大加快了我的训练时间，所以我更愿意保持这种方式，而不是回到分桶模型和/或将训练数据中的所有序列填充到相同的长度。我正在使用 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，但由于输入形状的未指定性质，我遇到了一个问题。根据tf.nn.sampled_softmax_loss 的文档，它需要为每个时间步单独提供输入。但是，我不能打电话，例如，

tf.unstack(target_labels, axis=1)

因为轴事先是未知的。有谁知道我该如何实现这个？有人会假设，由于 dynamic_rnn 和tf.losses.sparse_softmax_cross_entropy似乎都没有问题，因此可以以某种方式使用采样的 softmax 损失来实现一种解决方法。在深入研究了源代码甚至模型存储库之后，我空手而归。任何帮助/建议将不胜感激。