TensorFlow中gather, gather_nd, scatter, scatter_nd用法淺析

tf.gather和gather_nd從params中收集數值，tf.scatter_nd 和 tf.scatter_nd_update用updates更新某一張量。嚴格上說，tf.gather_nd和tf.scatter_nd_update互為逆操作。

1. 已知數值的位置，從張量中提取數值：tf.gather, tf.gather_nd

tf.gather indices每個元素（標量）是params某個axis的索引，tf.gather_nd 中indices最后一個階對應於索引值。

tf.gather函數

函數原型

gather(
    params,
    indices,
    validate_indices=None,
    name=None,
    axis=0
)

參數：

params是要查找的張量，indices是要查找值的索引（int32或int64），axis是查找軸，name是操作名。

如果indices是標量，$output[a_0,...,a_n,b_0,...,b_n] = params[a_0,...a_n,indices,b_0,...,b_n]$

如果indices是向量，$output[a_0,...,a_n,i,b_0,...,b_n] = params[a_0,...a_n,indices[i],b_0,...,b_n]$

如果indices是高階張量，$output[a_0,...,a_n,i,...,j,b_0,...,b_n] = params[a_0,...a_n,indices[i,...,j],b_0,...,b_n]$

返回值：

該函數返回值類型與params相同，具體值是從params中收集過來的，形狀為$params.shape[:axis]+indices.shape+params.shape[axis+1:]$。

 

tf.gather_nd函數

函數原型

gather_nd(
    params,
    indices,
    name=None
)

indices是K階張量，包含K-1階的索引值。它最后一階是索引，最后一階維度必須小於等於params的秩。indices最后一階的維數等於params的秩時，我們得到params的某些元素；indices最后一階的維數小於params的秩時，我們得到params的切片。例如$output[i_0,...,i_{K-2}]=params[indices[i_0,...i_{K-2}]]$。

輸出張量的形狀由indices的K-1階和params索引到的形狀拼接而成，形狀為$indices.shape[:-1]+params.shape[indices.shape[-1]:]$。

參數：

params：被收集的張量。

indices：索引張量。必須是以下類型之一：int32，int64。

name：操作的名稱(可選)。

返回值：

該函數返回一個張量.與params具有相同的類型。張量值從indices所給定的索引中收集，並且具有這樣的形狀：$indices.shape[:-1]+params.shape[indices.shape[-1]:]$

 

 2. 已知賦值的位置，向張量賦值：tf.scatter_nd, tf.scatter_nd_update

tf.scatter_nd對零張量進行賦值，tf.scatter_nd_update對已有可變的張量進行賦值。

tf.scatter_nd函數

scatter_nd(
    indices,
    updates,
    shape,
    name=None
)

創建一個形狀為shape的零張量，將updates賦值到indices指定的位置。

indices是整數張量，最內部維度對應於索引。indices最后一維的維數應不大於輸出張量shape的秩：$indices.shape[-1]<=shape.rank$。

如果indices.shape[-1] = shape.rank，那么indices直接對應到新張量的單個元素。如果indices.shape[-1] < shape.rank，那么indices中每個元素對新張量做切片操作。updates的形狀為$indices.shape[:-1]+shape[indices.shape[-1]:]$。

• 如果我們要把形狀為（4，）的updates賦值給形狀為（8，）的零張量，如下圖所示。

 我們需要這樣做

indices = tf.constant([[4], [3], [1], [7]])
updates = tf.constant([9, 10, 11, 12])
shape = tf.constant([8])
scatter = tf.scatter_nd(indices, updates, shape)
with tf.Session() as sess:
    print(sess.run(scatter))

我們得到這樣的張量

[0, 11, 0, 10, 9, 0, 0, 12]

上面代碼中，indices的形狀是（4,1），updates的形狀是（4，），shape的形狀是（8，）。

updates的形狀：$indices.shape[:-1]+shape[indices.shape[-1]:]=(4,)+(,)=(4,)$

• 如果我們要在三階張量中插入兩個切片，如下圖所示，則應該像下面代碼里所說的那樣做：

indices = tf.constant([[0], [2]])
updates = tf.constant([[[5, 5, 5, 5], [6, 6, 6, 6],
   [7, 7, 7, 7], [8, 8, 8, 8]],
   [[5, 5, 5, 5], [6, 6, 6, 6],
   [7, 7, 7, 7], [8, 8, 8, 8]]])
shape = tf.constant([4, 4, 4])
scatter = tf.scatter_nd(indices, updates, shape)
with tf.Session() as sess:
    print(sess.run(scatter))

indices的形狀是（2,1），updates的形狀是（2,4,4），shape的形狀是（4,4,4）。

updates的形狀：$indices.shape[:-1]+shape[indices.shape[-1]:]=(2,)+(4,4)=(2,4,4)$

我是這樣理解，indices的非最后一維的張量的形狀和updates張量是一樣的，用於和updates張量對齊，里邊indices最后一維的取值決定了updates對齊的這一個張量映射到shape(output)中的指向；而indices最后一維用於索引shape（output）這個張量，從圖中可以看出索引由外及內的順序。indices最后一維的維數不會超過shape(output)這個張量的秩。

對於張量的秩，可以根據張量的括號層數大致判斷，標量的秩是0，向量的秩是1，矩陣的秩是2，立方體的秩是3。我對張量的大致理解是對應不同組基向量的組合上的相應權重，更具體的判斷可以參考如何理解張量和張量的秩。

我們會得到下面的張量：

[[[5, 5, 5, 5], [6, 6, 6, 6], [7, 7, 7, 7], [8, 8, 8, 8]],
 [[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]],
 [[5, 5, 5, 5], [6, 6, 6, 6], [7, 7, 7, 7], [8, 8, 8, 8]],
 [[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]]]

函數參數

indices：Tensor；必須是以下類型之一：int32,int64；索引值張量。

updates：Tensor；分散到輸出的更新。

shape：Tensor；必須與indices具有相同的類型；1-d；得到的張量的形狀。

name：操作的名稱(可選)。

返回值

此函數返回一個Tensor,它與updates有相同的類型；一個有shape形狀的新張量，初始化值為0，部分值根據indices用updates進行更新。

 

tf.scatter_nd_update函數

函數原型

scatter_nd_update(
    ref,
    indices,
    updates,
    use_locking=True,
    name=None
)

scatter_nd_update也是把updates里面的值根據indices賦值到另外一個張量中，與scatter_nd不同的是，它是賦值到ref。

ref是秩為P的張量，indices是秩為Q的張量。

indices是整數類型的張量，必須具有這樣的形狀$[d_0,d_1,...,d_{Q-2},K], 0<K<=P$。

indices最內部的維度對應於ref的某個元素或切片。

updates的形狀是$[d_0,...,d_{Q-2},ref.shape[K],...,ref.shape[P-1]]$，是秩為Q-1+P-K的張量。

• 如果我們想要把（4，）的向量賦值到（8，）的ref中，我們可以像下面這樣操作。

ref = tf.Variable([1, 2, 3, 4, 5, 6, 7, 8])
indices = tf.constant([[4], [3], [1] ,[7]])
updates = tf.constant([9, 10, 11, 12])
update = tf.scatter_nd_update(ref, indices, updates)
with tf.Session() as sess:
    print sess.run(update)

我們可以得到這樣的ref

[1, 11, 3, 10, 9, 6, 7, 12]

函數參數

ref：一個可變的Tensor。

indices：一個 int32 或 int64 Tensor；一個對ref進行索引的張量.

updates：一個Tensor.必須與ref具有相同的類型；更新值張量.

use_locking：可選的bool；如果為True,則賦值將受鎖定的保護；否則行為是不確定的,但可能表現出較少的爭用.

name：操作的名稱(可選).

返回值：

經過更新的ref。

 

參考：

1. https://www.zhangshengrong.com/p/nDa9j05yNj/

2. https://www.tensorflow.org/api_docs/python/tf/scatter_nd