如何在TensorFlow中移除张量指定索引范围的数据？

原始张量定义 #

首先我们有通过tf.constant创建的原始张量：

import tensorflow as tf

params = tf.constant([
    1.3, 7, 1, 0.5, -2,
    3, -0.033, 0.9, -6.3, 4.1,
    9, 5, 0.25, -6, 0.2])

该张量的输出为：

<tf.Tensor: shape=(15,), dtype=float32, numpy=
array([ 1.3  ,  7.   ,  1.   ,  0.5  , -2.   ,  3.   , -0.033,  0.9  ,
       -6.3  ,  4.1  ,  9.   ,  5.   ,  0.25 , -6.   ,  0.2  ],
      dtype=float32)>

需求说明 #

需要移除以下元素：

• 第一个元素（值1.3，索引0）
• 索引4至6的元素（对应值-2、3、-0.033）
• 索引12及以后的元素（对应值0.25、-6、0.2）

期望得到的输出张量：

<tf.Tensor: shape=(8,), dtype=float32, numpy=
array([  7.   ,  1.   ,  0.5  ,   0.9  ,
       -6.3  ,  4.1  ,  9.   ,  5.],
      dtype=float32)>

实现方案 #

完全可以实现这个需求，我们可以通过筛选保留的索引来提取目标元素，具体代码如下：

# 定义需要保留的索引：索引1-3、7-11
keep_indices = tf.concat([tf.range(1, 4), tf.range(7, 12)], axis=0)
# 提取指定索引的元素
result = tf.gather(params, keep_indices)

print(result)

代码解释 #

1. tf.range(1, 4)生成索引1、2、3，对应原始张量中的7、1、0.5
2. tf.range(7, 12)生成索引7、8、9、10、11，对应原始张量中的0.9、-6.3、4.1、9、5
3. tf.concat将两组索引合并成一个一维索引张量
4. tf.gather根据合并后的索引从原始张量中提取元素，得到目标结果

运行上述代码后，输出结果与期望完全一致：

<tf.Tensor: shape=(8,), dtype=float32, numpy=
array([ 7. ,  1. ,  0.5,  0.9, -6.3,  4.1,  9. ,  5. ], dtype=float32)>

内容的提问来源于stack exchange，提问作者CA Khan