python - 使用Tensorflow SparseTensors进行有效的 bool 掩蔽

因此，我想掩盖SparseTensor的整个行。使用tf.boolean_mask可以很容易地做到这一点，但是SparseTensor并没有等效项。目前，对于我来说，可能的事情是只遍历SparseTensor.indices中的所有索引，并过滤掉所有非掩码行的索引，例如：

masked_indices = list(filter(lambda index: masked_rows[index[0]], indices))

其中masked_rows是一维数组，用于确定该索引处的行是否被屏蔽。

但是，这确实很慢，因为我的SparseTensor相当大（它有90k索引，但是会越来越大）。在我什至对过滤的索引应用SparseTensor.mask之前，在单个数据点上花费相当多的时间。这种方法的另一个缺点是，它实际上也不会删除所有行（尽管在我看来，全零的行也一样）。

有没有更好的方法来逐行屏蔽SparseTensor，还是最好的方法？

最佳答案

您可以这样做：

import tensorflow as tf

def boolean_mask_sparse_1d(sparse_tensor, mask, axis=0):  # mask is assumed to be 1D
    mask = tf.convert_to_tensor(mask)
    ind = sparse_tensor.indices[:, axis]
    mask_sp = tf.gather(mask, ind)
    new_size = tf.math.count_nonzero(mask)
    new_shape = tf.concat([sparse_tensor.shape[:axis], [new_size],
                           sparse_tensor.shape[axis + 1:]], axis=0)
    new_shape = tf.dtypes.cast(new_shape, tf.int64)
    mask_count = tf.cumsum(tf.dtypes.cast(mask, tf.int64), exclusive=True)
    masked_idx = tf.boolean_mask(sparse_tensor.indices, mask_sp)
    new_idx_axis = tf.gather(mask_count, masked_idx[:, axis])
    new_idx = tf.concat([masked_idx[:, :axis],
                         tf.expand_dims(new_idx_axis, 1),
                         masked_idx[:, axis + 1:]], axis=1)
    new_values = tf.boolean_mask(sparse_tensor.values, mask_sp)
    return tf.SparseTensor(new_idx, new_values, new_shape)

# Test
sp = tf.SparseTensor([[1], [3], [4], [6]], [1, 2, 3, 4], [7])
mask = tf.constant([True, False, True, True, False, False, True])
out = boolean_mask_sparse_1d(sp, mask)
print(out.indices.numpy())
# [[2]
#  [3]]
print(out.values.numpy())
# [2 4]
print(out.shape)
# (4,)

关于python - 使用Tensorflow SparseTensors进行有效的 bool 掩蔽，我们在Stack Overflow上找到一个类似的问题：https://stackoverflow.com/questions/57998859/