tensorflow2.1的索引与切片

1、常规索引

直接索引

import tensorflow as tf

a = tf.ones([1, 5, 5, 3])

a[0][0]
Out[5]:
<tf.Tensor: shape=(5, 3), dtype=float32, numpy=
array([[1., 1., 1.],
       [1., 1., 1.],
       [1., 1., 1.],
       [1., 1., 1.],
       [1., 1., 1.]], dtype=float32)>

a[0][0][0]
Out[6]: <tf.Tensor: shape=(3,), dtype=float32, numpy=array([1., 1., 1.], dtype=float32)>

a[0][0][0][2]
Out[7]: <tf.Tensor: shape=(), dtype=float32, numpy=1.0>

numpy风格的索引

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

a[1].shape
Out[11]: TensorShape([28, 28, 3])

a[1, 2].shape
Out[12]: TensorShape([28, 3])

a[1, 2, 3].shape
Out[13]: TensorShape([3])

a[1, 2, 3, 2].shape
Out[14]: TensorShape([])

使用start:end进行切片和使用 start: end :step 隔空采样

注:包含start不包含end，即，[start:end)

a = tf.range(10)

a[-1:]
Out[18]: <tf.Tensor: shape=(1,), dtype=int32, numpy=array([9])>

a[-2:]
Out[19]: <tf.Tensor: shape=(2,), dtype=int32, numpy=array([8, 9])>

a[:2]
Out[20]: <tf.Tensor: shape=(2,), dtype=int32, numpy=array([0, 1])>

a[:-1]
Out[21]: <tf.Tensor: shape=(9,), dtype=int32, numpy=array([0, 1, 2, 3, 4, 5, 6, 7, 8])>


#### 隔空采样
a.shape
Out[23]: TensorShape([4, 28, 28, 3])

a[:2, :, :, :].shape
Out[25]: TensorShape([2, 28, 28, 3])

a[:, 0:28:2, 0:28:2, :].shape
Out[26]: TensorShape([4, 14, 14, 3])

a[:, :14, :14, :].shape
Out[28]: TensorShape([4, 14, 14, 3])

a[:, 14:, 14:, :].shape
Out[29]: TensorShape([4, 14, 14, 3])

a[:, ::2, ::2, :].shape
Out[31]: TensorShape([4, 14, 14, 3])

实现一个倒序的功能

a = tf.range(4)

a[::-1]
Out[34]: <tf.Tensor: shape=(4,), dtype=int32, numpy=array([3, 2, 1, 0])>

a[::-1]
Out[35]: <tf.Tensor: shape=(4,), dtype=int32, numpy=array([3, 2, 1, 0])>

a[::-2]
Out[36]: <tf.Tensor: shape=(2,), dtype=int32, numpy=array([3, 1])>

a[2::-2]
Out[37]: <tf.Tensor: shape=(2,), dtype=int32, numpy=array([2, 0])>

省略号（…）替代任意长

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

a[0].shape
Out[39]: TensorShape([4, 28, 28, 3])

a[0, :, :, :, :].shape
Out[41]: TensorShape([4, 28, 28, 3])

a[0, ...].shape
Out[43]: TensorShape([4, 28, 28, 3])

a[:,:, :, :, 0].shape
Out[44]: TensorShape([2, 4, 28, 28])

a[..., 0].shape
Out[45]: TensorShape([2, 4, 28, 28])

a[0, ..., 2].shape
Out[46]: TensorShape([4, 28, 28])

a[1, 0, ..., 0].shape
Out[47]: TensorShape([28, 28])

2、选择索引（tf.gather、tf.gather_nd、tf.boolean_mask）

tf.gather

axis - 代表维度

indices - 索引号

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

tf.gather(a, axis = 0, indices = [2, 3]).shape
Out[49]: TensorShape([2, 35, 8])

a[2:4].shape
Out[50]: TensorShape([2, 35, 8])

tf.gather(a, axis = 0, indices = [2, 1, 3, 0]).shape
Out[52]: TensorShape([4, 35, 8])

tf.gather(a, axis = 1, indices = [2,3 , 7, 9, 16]).shape
Out[54]: TensorShape([4, 5, 8])

tf.gather(a, axis = 2, indices = [2,3 , 7]).shape
Out[55]: TensorShape([4, 35, 3])

• 使用gather实现多个维度的同时的选择

假设有4个班级（class），35个学生（student），8门课程，[4, 35, 8],这个三维数据，现在想取其中3个学生(这3个学生是[4， 7， 8]号学生)的3门课程（这3门课程是[4, 7, 3]）。

可以这样写：

aa = tf.gather(a, axis = 1, indices = [4, 7, 8])
aaa = tf.gather(aa, axis = 2, indices = [4, 7, 3])

tf.gather_nd

假设有4个班级（class），35个学生（student），8门课程，[4, 35, 8],这个三维数据，现在想取2号学生的第0门课程， 3号学生的第4门课程， 8号学生的第2门课程[2->0, 3->4, 8->2],也即学生和课程都是对应 抽取的。

a.shape
Out[58]: TensorShape([4, 35, 8])

tf.gather_nd(a, [0]).shape
Out[59]: TensorShape([35, 8])

tf.gather_nd(a, [0, 1]).shape
Out[60]: TensorShape([8])

tf.gather_nd(a, [0, 1, 2]).shape
Out[61]: TensorShape([])

tf.gather_nd(a, [[0, 1, 2]]).shape
Out[62]: TensorShape([1])

tf.gather_nd(a, [[0, 0], [1, 1]]).shape
Out[64]: TensorShape([2, 8])

tf.gather_nd(a, [[0, 0], [1, 1], [2, 2]]).shape
Out[65]: TensorShape([3, 8])

tf.gather_nd(a, [[0, 0, 0], [1, 1, 1], [2, 2, 2]]).shape
Out[66]: TensorShape([3])

tf.gather_nd(a, [[[0, 0, 0], [1, 1, 1], [2, 2, 2]]]).shape
Out[67]: TensorShape([1, 3])

tf.boolean_mask

通过mask过滤元素, 默认轴为0

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

tf.boolean_mask(a, mask = [True, True, False, False]).shape
Out[69]: TensorShape([2, 28, 28, 3])

tf.boolean_mask(a, mask = [True, True, False], axis = 3).shape
Out[70]: TensorShape([4, 28, 28, 2])

# 多维组合
a = tf.ones([2, 3, 4])
# 注：[[True, False, False],
#      [False, True, True]]是一个2行3列的mask，对应[2, 3, 4]的前两维[2, 3]数据。
tf.boolean_mask(a, mask = [[True, False, False], [False, True, True]])
Out[72]:
<tf.Tensor: shape=(3, 4), dtype=float32, numpy=
array([[1., 1., 1., 1.],
       [1., 1., 1., 1.],
       [1., 1., 1., 1.]], dtype=float32)>