Tensorflow稀疏张量逐行掩码

Question

我有一个tf稀疏张量，和一个张量稠密形式的行掩码。

。

tf.SparseTensor(
indices=[[0, 0], [0, 1], [0, 2], [1, 0], [1, 1], [1, 2]],
values=tf.constant([3, 40, 9, 5, 6, 4], tf.int64),
dense_shape=[2, 10]
)

我有一个掩码[True, False]，我想保持第一行作为密集形式，不转换稀疏张量:

tf.SparseTensor(
indices=[[0, 0], [0, 1], [0, 2]],
values=tf.constant([3, 40, 9], tf.int64),
dense_shape=[1, 10]
)

如何使用掩码直接过滤稀疏张量?

Answer 1

也许可以尝试使用boolean_mask和tf.sparse.slice:

import tensorflow as tf
x = tf.SparseTensor(
indices=[[0, 0], [0, 1], [0, 2], [1, 0], [1, 1], [1, 2]],
values=tf.constant([3, 40, 9, 5, 6, 4], tf.int64),
dense_shape=[2, 10])
mask = tf.constant([True, False])
indices = x.indices
start = tf.boolean_mask(indices, tf.where(tf.greater(indices[:, 0], 0), mask[1], mask[0]), axis=0)
result = tf.sparse.slice(x, start = start[0,:], size = tf.cast(tf.shape(tf.boolean_mask(tf.zeros(x.dense_shape), mask, axis=0)), dtype=tf.int64))
print(result)

SparseTensor(indices=tf.Tensor(
[[0 0]
[0 1]
[0 2]], shape=(3, 2), dtype=int64), values=tf.Tensor([ 3 40  9], shape=(3,), dtype=int64), dense_shape=tf.Tensor([ 1 10], shape=(2,), dtype=int64))

或者使用tf.sparse.retain，您可以直接在稀疏张量上使用布尔掩码:

import tensorflow as tf
x = tf.SparseTensor(
indices=[[0, 0], [0, 1], [0, 2], [1, 0], [1, 1], [1, 2]],
values=tf.constant([3, 40, 9, 5, 6, 4], tf.int64),
dense_shape=[2, 10])
mask = tf.constant([True, False])
mask = tf.where(tf.greater(x.indices[:, 0], 0), mask[1], mask[0])
result = tf.sparse.retain(x, mask)
print(result)

但是，这种方法不会改变dense_shape:

SparseTensor(indices=tf.Tensor(
[[0 0]
[0 1]
[0 2]], shape=(3, 2), dtype=int64), values=tf.Tensor([ 3 40  9], shape=(3,), dtype=int64), dense_shape=tf.Tensor([ 2 10], shape=(2,), dtype=int64))

但是你可以在result上使用tf.sparse.reduce_sum来得到具有正确dense_shape的稀疏张量:

tf.sparse.reduce_sum(result, axis=0, keepdims=True, output_is_sparse=True)