Tensorflow2.0之根据索引号收集数据

Tensorflow之根据索引号收集数据

tf.gather

tf.gather 可以实现根据索引号收集数据的目的。考虑班级成绩册的例子，共有4个班，每个班35个学生，8门学科，保存成绩册的张量shape为[4,35,8]。

x = tf.random.uniform([4,35,8],maxval=100,dtype=tf.int32)

现在需要收集1-2班的成绩册，可以给定班级索引号：[0,1]/[:2]，班级的维度为axis=0:

In [38]:tf.gather(x,[0,1],axis=0) # 在班级维度收集第1-2 号班级成绩册
Out[38]:<tf.Tensor: id=83, shape=(2, 35, 8), dtype=int32, numpy=
array([[[43, 10, 93, 85, 75, 87, 28, 19],
[52, 17, 44, 88, 82, 54, 16, 65],
[98, 26, 1, 47, 59, 3, 59, 70],…

实际上，对于上述需求，通过切片𝑥[: 2]可以更加方便地实现。但是对于不规则的索引方式，比如，需要抽查所有班级的第1,4,9,12,13,27 号同学的成绩，则切片方式实现起来非常麻烦，而tf.gather 则是针对于此需求设计的，使用起来非常方便：

In [39]: # 收集第1,4,9,12,13,27 号同学成绩
tf.gather(x,[0,3,8,11,12,26],axis=1)
Out[39]:<tf.Tensor: id=87, shape=(4, 6, 8), dtype=int32, numpy=
array([[[43, 10, 93, 85, 75, 87, 28, 19],
[74, 11, 25, 64, 84, 89, 79, 85],…

如果需要收集所有同学的第3，5 等科目的成绩，则可以：

In [40]:tf.gather(x,[2,4],axis=2) # 第3，5 科目的成绩
Out[40]:<tf.Tensor: id=91, shape=(4, 35, 2), dtype=int32, numpy=
array([[[93, 75],
[44, 82],[ 1, 59],…

可以看到，tf.gather 非常适合索引没有规则的场合，其中索引号可以乱序排列，此时收集的数据也是对应顺序:

In [41]:a=tf.range(8)
a=tf.reshape(a,[4,2]) # 生成张量a
Out[41]:<tf.Tensor: id=115, shape=(4, 2), dtype=int32, numpy=
array([[0, 1],
[2, 3],
[4, 5],
[6, 7]])>
In [42]:tf.gather(a,[3,1,0,2],axis=0) # 收集第4,2,1,3 号元素
Out[42]:<tf.Tensor: id=119, shape=(4, 2), dtype=int32, numpy=
array([[6, 7],
[2, 3],
[0, 1],
[4, 5]])>

我们将问题变得复杂一点：如果希望抽查第[2,3]班级的第[3,4,6,27]号同学的科目成绩，则可以通过组合多个tf.gather 实现。首先抽出第[2,3]班级：

In [43]:students=tf.gather(x,[1,2],axis=0) # 收集第2,3 号班级
Out[43]:<tf.Tensor: id=227, shape=(2, 35, 8), dtype=int32, numpy=
array([[[ 0, 62, 99, 7, 66, 56, 95, 98],…

再从这2 个班级的同学中提取对应学生成绩：

In [44]:tf.gather(students,[2,3,5,26],axis=1) # 收集第3,4,6,27 号同学
Out[44]:<tf.Tensor: id=231, shape=(2, 4, 8), dtype=int32, numpy=
array([[[69, 67, 93, 2, 31, 5, 66, 65], …

此时得到这2 个班级4 个同学的成绩，shape 为[2,4,8]。

我们再将问题进一步复杂：这次我们希望抽查第2 个班级的第2 个同学的所有科目，
第3 个班级的第3 个同学的所有科目，第4 个班级的第4 个同学的所有科目。那么怎么实现呢？

可以通过笨方式一个一个的手动提取：首先提取第一个采样点的数据：𝑥[1,1],可得到
8 门科目的数据向量：

In [45]: x[1,1] # 收集第2 个班级的第2 个同学
Out[45]:<tf.Tensor: id=236, shape=(8,), dtype=int32, numpy=array([45, 34,
99, 17, 3, 1, 43, 86])>

再串行提取第二个采样点的数据：𝑥[2,2]，和第三个采样点的数据𝑥[3,3]，最后通过stack方式合并采样结果：

In [46]: tf.stack([x[1,1],x[2,2],x[3,3]],axis=0)
Out[46]:<tf.Tensor: id=250, shape=(3, 8), dtype=int32, numpy=
array([[45, 34, 99, 17, 3, 1, 43, 86],
[11, 25, 84, 95, 97, 95, 69, 69],
[ 0, 89, 52, 29, 76, 7, 2, 98]])>

这种方法也能正确的得到shape 为[3,8]的结果，其中3 表示采样点的个数，4 表示每个采样点的数据。但是它最大的问题在于手动串行方式执行采样，计算效率极低。有没有更好的方式实现呢？这就是下面要介绍的tf.gather_nd 的功能。

tf.gather_nd

通过 tf.gather_nd，可以通过指定每次采样的坐标来实现采样多个点的目的。回到上面的挑战，我们希望抽查第2 个班级的第2 个同学的所有科目，第3 个班级的第3 个同学的所有科目，第4 个班级的第4 个同学的所有科目。那么这3 个采样点的索引坐标可以记为:[1,1], [2,2], [3,3]，我们将这个采样方案合并为一个List 参数：[[1,1], [2,2], [3,3]]，通过tf.gather_nd 实现如下：

In [47]: # 根据多维度坐标收集数据
tf.gather_nd(x,[[1,1],[2,2],[3,3]])
Out[47]:<tf.Tensor: id=256, shape=(3, 8), dtype=int32, numpy=
array([[45, 34, 99, 17, 3, 1, 43, 86],
[11, 25, 84, 95, 97, 95, 69, 69],
[ 0, 89, 52, 29, 76, 7, 2, 98]])>

可以看到，结果与串行采样完全一致，实现更加简洁，计算效率大大提升。一般地，在使用tf.gather_nd 采样多个样本时，如果希望采样第i 号班级，第j 个学生，第k 门科目的成绩，则可以表达为[. . . , [𝑖, 𝑗, 𝑘], . . . ]，外层的括号长度为采样样本的个数，内层列表包含了每个采样点的索引坐标：

In [48]: # 根据多维度坐标收集数据
tf.gather_nd(x,[[1,1,2],[2,2,3],[3,3,4]])
Out[48]:<tf.Tensor: id=259, shape=(3,), dtype=int32, numpy=array([99, 95,
76])>

上述代码中，我们抽出了班级1，学生1 的科目2；班级2，学生2 的科目3；班级3，学生3 的科目4 的成绩，共有3 个成绩数据，结果汇总为一个shape 为[3]的张量。