TensorFlow张量提取与赋值：tf.gather系列与tf.scatter_nd系列详解

TensorFlow张量提取与赋值概述

在TensorFlow中，tf.gather和tf.gather_nd用于从参数张量params中收集特定数值，而tf.scatter_nd和tf.scatter_nd_update则用于将updates中的值更新到目标张量的指定位置。严格来说，tf.gather_nd与tf.scatter_nd_update互为逆操作。

当已知数值在张量中的具体位置时，可以使用tf.gather或tf.gather_nd来提取这些数值。

在tf.gather中，indices的每个元素（标量）代表params中某个axis的索引；而在tf.gather_nd中，indices的最后一阶对应于具体的索引值。

tf.gather系列：张量数值提取

tf.gather函数详解

函数原型如下所示。

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

其中，params是待查找的源张量，indices是包含待查找值索引的张量（类型为int32或int64），axis指定查找的轴，name为可选的操作名称。

当indices为标量时：

当indices为向量时：

当indices为高阶张量时：

关于返回值：

该函数的返回值类型与params保持一致，具体数值是从params中根据indices收集而来，其输出形状由indices的形状和params中对应切片决定。

tf.gather_nd：高阶索引提取

tf.gather_nd函数详解

函数原型如下所示。

gather_nd(
 params,
 indices,
 name=None
)

indices是一个K阶张量，包含K-1阶的索引值。其最后一阶表示具体的索引，且最后一阶的维度必须小于等于params的秩。当indices最后一阶的维数等于params的秩时，提取的是params的某些特定元素；当indices最后一阶的维数小于params的秩时，提取的是params的切片。

输出张量的形状由indices的K-1阶形状与params被索引部分的形状拼接而成，具体计算方式如下。

indices.shape[:-1] + params.shape[indices.shape[-1]:]

参数说明如下：

params：被收集数值的源张量。

indices：索引张量，类型必须为int32或int64。

name：操作的名称，可选参数。

关于返回值：

该函数返回一个张量，其类型与params相同。张量的值从indices指定的索引处收集，其形状由indices的K-1阶和params被索引部分的形状拼接而成。

tf.scatter_nd系列：张量数值赋值

当已知赋值位置时，可以使用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[-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，）。

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）。

indices.shape[:-1]+shape[indices.shape[-1]:]=(2,)+(4,4)=(2,4,4)

最终得到的张量结果为：

[[[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]]]

tf.scatter_nd函数参数说明：

indices：Tensor类型，必须为int32或int64，表示索引值张量。

updates：Tensor类型，表示分散到输出中的更新值。

shape：Tensor类型，必须与indices具有相同的类型，为一维张量，表示输出张量的形状。

name：操作的名称，可选参数。

关于返回值：

此函数返回一个Tensor，其类型与updates相同。这是一个具有shape形状的新张量，初始值为0，部分值根据indices使用updates进行更新。

tf.scatter_nd_update：更新可变张量

tf.scatter_nd_update函数详解

函数原型如下所示。

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

tf.scatter_nd_update同样根据indices将updates中的值赋值到目标张量，但与scatter_nd不同，它是直接赋值到ref变量中。

ref是秩为P的张量，indices是秩为Q的张量。

indices是整数类型的张量，其形状需满足特定结构。

indices最内部的维度对应于ref的某个元素或切片。

updates的形状为，是秩为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]

tf.scatter_nd_update函数参数说明：

ref：一个可变的Tensor。

indices：一个int32或int64 Tensor，用于对ref进行索引。

updates：一个Tensor，必须与ref具有相同的类型，表示更新值张量。

use_locking：可选的bool类型参数。如果为True，则赋值操作将受锁保护；否则行为不确定，但可能表现出较少的竞争。

name：操作的名称，可选参数。

关于返回值：

返回经过更新后的ref。

以上关于TensorFlow中张量提取值和赋值的介绍到此结束，希望能为读者提供参考。