python - 用 tensorflow 中的条件减少总和

Question

我得到了一个带有随机行的二维张量。应用 tf.math.greater() 和 tf.cast(tf.int32) 后，我留下了一个包含 0 和 1 的张量。我现在想将归约总和应用到该矩阵上，但有一个条件:如果至少有一个 1 求和，并且后面有一个 0，我也想删除所有后面的 1，这意味着 1 0 1 应该导致1 而不是 2。

我尝试使用 tf.scan() 解决问题，但我无法想出一个能够处理起始 0 的函数，因为该行可能看起来像: 0 0 0 1 0 1一个想法是将矩阵的下半部分设置为 1(因为我知道对角线剩下的所有内容都将始终为 0)，然后运行类似 tf.scan() 的函数来过滤掉点(请参阅下面的代码和错误消息)。

设 z 为 tf.cast 之后的矩阵。

helper = tf.matrix_band_part(tf.ones_like(z), -1, 0)
z = tf.math.logical_or(tf.cast(z, tf.bool), tf.cast(helper,tf.bool))
z = tf.cast(z, tf.int32)
z = tf.scan(lambda a, x: x if a == 1 else 0 ,z)

结果:

ValueError:值 ([]) 的形状不兼容，预期 ([5])

Answer 1

IIUC，这是一种无需扫描或循环即可完成所需操作的方法。它可能有点复杂，并且实际上迭代了两次列(一次 cumsum 和一次 cumprod)，但作为矢量化操作，我认为它可能会更快。代码是 TF 2.x，但在 TF 1.x 中运行相同(显然除了最后一行)。

import tensorflow as tf

# Example data
a = tf.constant([[0, 0, 0, 0],
                 [1, 0, 0, 0],
                 [0, 1, 1, 0],
                 [0, 1, 0, 1],
                 [1, 1, 1, 0],
                 [1, 1, 0, 1],
                 [0, 1, 1, 1],
                 [1, 1, 1, 1]])
# Cumsum columns
c = tf.math.cumsum(a, axis=1)
# Column-wise differences
diffs = tf.concat([tf.ones([tf.shape(c)[0], 1], c.dtype), c[:, 1:] - c[:, :-1]], axis=1)
# Find point where we should not sum anymore (cumsum is not zero and difference is zero)
cutoff = tf.equal(a, 0) & tf.not_equal(c, 0)
# Make mask
mask = tf.math.cumprod(tf.dtypes.cast(~cutoff, tf.uint8), axis=1)
# Compute result
result = tf.reduce_max(c * tf.dtypes.cast(mask, c.dtype), axis=1)
print(result.numpy())
# [0 1 2 1 3 2 3 4]