tensorflow中的gather、gather_nd、boolean_mask

1. gather

import tensorflow as tf

a = tf.random.normal([4, 38, 8])
print(a.shape)

a_1 = tf.gather(a, axis=0, indices=[3, 1]) # 第一维度的3和1，这里的indices是不用限制顺序的
print(a_1.shape)

a_2 = tf.gather(a, axis=1, indices=[4, 5, 3, 1]) #第二个维度的4,5,3,1, shape[4,4,8]
print(a_2.shape)

a_3 = tf.gather(a, axis=2, indices=[1, 7]) #取第三个维度的1和7， shape[4,38,2]
print(a_3.shape)

#取第二个维度2的7,9,20,10维的前3个数据
a_4 = tf.gather(a[1], axis=0, indices=[7,9,20,10])
print(a_4.shape) #shape[4,8]
a_5 = tf.gather(a_4, axis=1, indices=[0,1,2])
print(a_5.shape) #shape[4,3]

其中注意的就是axis就是哪个维度，indices就是这个维度上取那几个。比如这样一个tensor [4,38,8]，取第二个维度11、13、20
参数就是这样的（输入的tensor,axis=1, indices = [11,13,20]）

2.gather_nd

import tensorflow as tf
# tf.gather_nd

a = tf.random.normal([4, 38, 8])

a_1 = tf.gather_nd(a, [0])
print(a_1.shape)  # 取第一个维度的0，shape[38, 8]

a_2 = tf.gather_nd(a, [0, 20]) # 第一个维度取0，第二个维度取20，shape[8]
print(a_2.shape)

a_3 = tf.gather_nd(a, [0, 21, 1]) # 第一个维度取0， 第二个取21，第三个取1 shape[]
print(a_3.shape)

a_4 = tf.gather_nd(a, [[0,21,1]])
print(a_4.shape)
print((a_4))


a_5 = tf.gather_nd(a,[[0],[1]]) # 取维度0的0,1 shape[2,38,8]
print(a_5.shape)

a_6 = tf.gather_nd(a, [[0,1],[1,1],[2,2]]) # shape[3,8]  从左向右看。如果[]中数字不足以覆盖后面的，就当后面的全取
print(a_6.shape)

a_7 = tf.gather_nd(a, [[0,0,0],[1,1,1],[2,2,2]]) # shape[3] 取得就是a[0,0,0],a[1,1,1],a[2,2,2]
print(a_7.shape)

a_8 = tf.gather_nd(a, [[[0,0,0],[1,1,1],[2,2,2]]]) # spape[1,3]
print(a_8.shape) #将倒数第二个[]看成一个整体

用法很简单主要是注意这里的a_7和a_8的区别，a_8多了一个[],代表它将取出的结果又当成一个新的tensor

3.boolean_mask

import tensorflow as tf

a = tf.random.normal([4,28,28,3])
print(a.shape)

a_1 = tf.boolean_mask(a,[True,False,False,True]) #默认axis=0,然后mask要和所选的维度一样大
print(a_1.shape)

a_2 = tf.boolean_mask(a, [True,False,False], axis=3) 

b = tf.random.normal([2,3,4])
b_1 = tf.boolean_mask(b, [[True,False,False],[False,True,False]])
print(b_1.shape)

如果使用时没有设置axis就进行默认推断，比如a_1中就默认为第0个维度
掩膜的大小要和自己想遮盖的范围一样大，可以通过axis设置自己的掩膜是用于那些维度。