python - 稀疏张量上的行或元素选择

Question

在 tensorflow 中，我们如何在稀疏张量中执行 tf.gather 或 tf.gather_nd？我们如何从稀疏张量中提取特定行或特定元素而不将其转换为密集张量？

Answer 1

这是一种可能的解决方案，尽管它在时间和内存方面仍然很昂贵，因此对于大型用例来说可能不可行:

import tensorflow as tf

def sparse_select_indices(sp_input, indices, axis=0):
    # Only necessary if indices may have non-unique elements
    indices, _ = tf.unique(indices)
    n_indices = tf.size(indices)
    # Only necessary if indices may not be sorted
    indices, _ = tf.math.top_k(indices, n_indices)
    indices = tf.reverse(indices, [0])
    # Get indices for the axis
    idx = sp_input.indices[:, axis]
    # Find where indices match the selection
    eq = tf.equal(tf.expand_dims(idx, 1), tf.cast(indices, tf.int64))
    # Mask for selected values
    sel = tf.reduce_any(eq, axis=1)
    # Selected values
    values_new = tf.boolean_mask(sp_input.values, sel, axis=0)
    # New index value for selected elements
    n_indices = tf.cast(n_indices, tf.int64)
    idx_new = tf.reduce_sum(tf.cast(eq, tf.int64) * tf.range(n_indices), axis=1)
    idx_new = tf.boolean_mask(idx_new, sel, axis=0)
    # New full indices tensor
    indices_new = tf.boolean_mask(sp_input.indices, sel, axis=0)
    indices_new = tf.concat([indices_new[:, :axis],
                             tf.expand_dims(idx_new, 1),
                             indices_new[:, axis + 1:]], axis=1)
    # New shape
    shape_new = tf.concat([sp_input.dense_shape[:axis],
                           [n_indices],
                           sp_input.dense_shape[axis + 1:]], axis=0)
    return tf.SparseTensor(indices_new, values_new, shape_new)

这是一个使用示例:

import tensorflow as tf

with tf.Session() as sess:
    # Input
    sp1 = tf.SparseTensor([[0, 1], [2, 3], [4, 5]], [10, 20, 30], [6, 7])
    print(sess.run(tf.sparse.to_dense(sp1)))
    # [[ 0 10  0  0  0  0  0]
    #  [ 0  0  0  0  0  0  0]
    #  [ 0  0  0 20  0  0  0]
    #  [ 0  0  0  0  0  0  0]
    #  [ 0  0  0  0  0 30  0]
    #  [ 0  0  0  0  0  0  0]]

    # Select rows 0, 1, 2
    sp2 = sparse_select_indices(sp1, [0, 1, 2])
    print(sess.run(tf.sparse.to_dense(sp2)))
    # [[ 0 10  0  0  0  0  0]
    #  [ 0  0  0  0  0  0  0]
    #  [ 0  0  0 20  0  0  0]]

    # Select columns 4, 5
    sp3 = sparse_select_indices(sp1, [4, 5], axis=1)
    print(sess.run(tf.sparse.to_dense(sp3)))
    # [[ 0  0]
    #  [ 0  0]
    #  [ 0  0]
    #  [ 0  0]
    #  [ 0 30]
    #  [ 0  0]]