2

我想计算批量张量的softmax_cross_entropy_with_logits。我有一批 logits 张量作为输入，但是我想在之前用布尔掩码屏蔽这个张量。布尔掩码也是一批掩码，每个掩码中可以有不同数量的True. 因此，将这个掩码应用于整个张量之后将不会变得密集。
尝试此方法可以使张量 ( tf.boolean_mask) 变平或创建一个参差不齐的 ( tf.ragged.boolean_mask)，这都会产生错误的结果或不适用于 softmax 函数。

所以基本上我想让以下代码工作：

# logits.shape = (batch, outputs), e.g. (512,8)
# mask.shape = (batch, valid), e.g. (512,8)
# expected result shape (512,)
one_hot_actions = tf.one_hot(x, logits.get_shape().as_list()[-1])
stopgradient = tf.stop_gradient(one_hot_actions)
return tf.nn.softmax_cross_entropy_with_logits_v2(
    logits=tf.boolean_mask(logits, mask),
    labels=tf.boolean_mask(stopgradient, mask))

但是tf.boolean_mask这样只产生一个值，而不是四个，并且tf.ragged.boolean_mask这个函数不起作用。

我尝试按行组合两个参差不齐的张量（第一个掩码 logits 行和第一个掩码标签行）并逐行计算 softmax。这不起作用，因为我使用的 tf.map_fn 不接受参差不齐的张量作为输入。我也用 SparseTensors 和张量列表（tf.split）尝试了同样的想法，但从来没有从中得到任何工作代码。

我必须解决这个问题的唯一想法是非常丑陋。tf.where用to替换所有掩码值，NaN然后在这些现在密集的张量上使用 map_fn。然后我可以再次屏蔽每一行以排除 NaN，现在可以再次逐行调用 softmax 函数。
编辑这是代码当前的样子：

stopgradient = tf.stop_gradient(one_hot_actions)
nan_logits = tf.where(mask, logits, float('NaN') + tf.zeros_like(self.logits))
nan_labels = tf.where(mask, stopgradient, float('NaN') + tf.zeros_like(stopgradient))
nan_lola = tf.stack([nan_logits, nan_labels], axis=1)
def fn(x):
    nan_lo = x[0]
    nan_la = x[1]
    masked_lo = tf.boolean_mask(nan_lo, tf.logical_not(tf.math.is_nan(nan_lo)))
    masked_la = tf.boolean_mask(nan_la, tf.logical_not(tf.math.is_nan(nan_la)))
    return tf.nn.softmax_cross_entropy_with_logits_v2(
        logits=masked_lo,
        labels=masked_la
    )
result = tf.map_fn(fn, nan_lola)
return result

这行得通，但速度很慢（我的训练时间几乎翻了一番）。

致那些感兴趣的人：我偶然发现了这个问题，试图掩盖强化学习中的有效行为。

你知道有什么方法可以做到这一点（更快）吗？你能用不影响softmax的值替换掩码值吗？谢谢！

4

0 回答 0