tf.gather,取指定維度多個索引的數據

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([01, 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,從每行取指定索引元素

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章