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)

Tensorflow 稀疏张量行掩码

Tensorflow sparse tensor row-wise mask

python

tensorflow