使用 Tensorflow SparseTensors 的高效布尔掩码
Efficient boolean masking with Tensorflow SparseTensors
所以,我想屏蔽掉 SparseTensor
的整行。使用 tf.boolean_mask
很容易做到这一点,但没有 SparseTensor
s 的等价物。目前,对我来说可能的是遍历 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,)
所以,我想屏蔽掉 SparseTensor
的整行。使用 tf.boolean_mask
很容易做到这一点,但没有 SparseTensor
s 的等价物。目前,对我来说可能的是遍历 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,)