张量流沿维度选择索引列表

Question

为了选择矩阵中的列列表，我正在执行以下操作：

sel = tf.concat([tf.slice(mat, [0, i], [-1, 1]) for i in list_columns],
                axis=1)

我想知道是否有更有效的方式

Answer 1

tf.gather将更加高效和简洁。让axis=1，然后您可以选择指定索引中的列。

mat = tf.constant(np.arange(12).reshape(2,6))
#[[ 0,  1,  2,  3,  4,  5],
# [ 6,  7,  8,  9, 10, 11]]
list_columns = [0,2,4]

res = tf.gather(mat, [0,2,4], axis=1)
#[[ 0,  2,  4],
# [ 6,  8, 10]]