在神经网络计算过程中,经常会遇到需要将矩阵中的某些元素取出并且单独进行计算的步骤(例如MLE,Attention等操作)。那么在 tensorflow 的 Variable 类型中如何做到这一点呢?
首先假设 Variable 是一个一维数组 A:
import numpy as np
import tensorflow as tf
a = np.array([1, 2, 3, 4, 5, 6, 7, 8])
A = tf.Variable(a)
我们把我们想取出的元素的索引存到 B 中,如果我们只想取出数组 A 中的某一个元素,则 B 的设定为:
b = np.array([3])
B = tf.placeholder(dtype=tf.int32, shape=[1])
由于我们的索引坐标只有一维,所以 shape=1。
取出元素然后组合成tensor C 的操作如下:
C = tf.gather_nd(A, B)
运行:
init = tf.global_variables_initializer()
with tf.Session() as sess:
init.run()
feed_dict = {B: b}
result = sess.run([C], feed_dict=feed_dict)
print result
得到:
[4]
如果我们想取出一维数组中的多个元素,则需要把每一个想取出的元素索引都单独放一行