tensorflow和numpy在数据处理上语法相似但又不完全一样,比如在numpy中想取指定维度的多个指定索引所指向的数据时,直接用一个列表保存索引就能直接取,比如:
# b的shape为[2, 3, 2]
b = np.array([[[1, 2], [2, 3], [3,4]], [[4, 5], [5,6], [6, 7]]])
a = [0, 1, 2]
# 假若要取b的第2个维度(从1算起)的以a为索引的数据,只需写成如下形式
b[:, a, :]
但是若b是tensor形式,则上述操作会报错!
TypeError: Only integers, slices (`:`), ellipsis (`...`), tf.newaxis (`None`) and scalar tf.int32/tf.int64 tensors are valid indices, got array([0, 1, 2])
而在tensor中则需要使用tf.gather方法来实现该功能。
tf.gather可以取指定维度任意索引的数据,举个例子:
b = np.array([[[1, 2], [2, 3], [3,4]], [[4, 5], [5,6], [6, 7]]])
b = tf.constant(b)
a = [0, 1, 2]
# 写法如下,效果和上面的numpy的效果相同
tf.gather(b, axis=1, indices=a)
这是tf.gather的用法,而想来大家也都听过tf.gather_nd,这个功能更加强大,它能对同一个维度的每一组数据分别取不同的索引值,比如取第1行的第二列+第二行的第三列。
关于gather_nd的用法见:tensorflow中tensor,从每行取指定索引元素