Question

我想在TensorFlow中做类似这段Numpy代码的事情：

a = np.zeros([5, 2])
idx = np.random.randint(0, 2, (5,))
row_idx = np.arange(5)
a[row_idx, idx] = row_idx

表示用另一个张量索引2D张量的所有行，然后为其分配张量。我对如何实现这一目标一无所知。

到目前为止，我在Tensorflow中可以做的事情如下

a = tf.Variable(tf.zeros((5, 2)))
idx = tf.constant([0, 1, 1, 0, 1])
row_idx = tf.range(5)
indices = tf.transpose([row_idx, idx])
r = tf.gather_nd(a, indices)
tf.assign(r, row_idx) # This line does not work

当我尝试执行此操作时，在最后一行出现以下错误：

AttributeError: 'Tensor' object has no attribute 'assign'

有没有解决的办法？必须有一些不错的方法来执行此操作，我不想在数据上使用for循环进行迭代，而是在每个元素的基础上手动进行分配。我知道，现在数组索引还不如Numpy的功能先进，但这应该还是可以实现的。

Answer 1

您经常尝试使用tf.scatter_nd_update完成操作。但是，多数情况下这不是正确的方法，您不需要一个变量，只需从原始张量生成的另一个张量并替换一些值即可。不幸的是，通常没有直接的方法可以做到这一点。如果您的原始张量实际上全为零，那么您可以简单地使用tf.scatter_nd：

import tensorflow as tf

idx = tf.constant([0, 1, 1, 0, 1])
row_idx = tf.range(5)
indices = tf.stack([row_idx, idx], axis=1)
a = tf.scatter_nd(indices, row_idx, (5, 2))
with tf.Session() as sess:
    print(sess.run(a))
# [[0 0]
#  [0 1]
#  [0 2]
#  [3 0]
#  [0 4]]

但是，如果初始张量不全为零，则更为复杂。一种方法是执行与上述相同的操作，然后为更新内容制作一个遮罩，然后根据遮罩在原始遮罩和更新遮罩之间进行选择：

import tensorflow as tf

a = tf.ones((5, 2), dtype=tf.int32)
idx = tf.constant([0, 1, 1, 0, 1])
row_idx = tf.range(5)
indices = tf.stack([row_idx, idx], axis=1)
a_update = tf.scatter_nd(indices, row_idx, (5, 2))
update_mask = tf.scatter_nd(indices, tf.ones_like(row_idx, dtype=tf.bool), (5, 2))
a = tf.where(update_mask, a_update, a)
with tf.Session() as sess:
    print(sess.run(a))
# [[0 1]
#  [1 1]
#  [1 2]
#  [3 1]
#  [1 4]]

Answer 2

我不知道以前的版本，但是在Tensorflow 2.1中，您可以使用tf.tensor_scatter_nd_update单行执行所需的操作。在您的代码示例中，您可以执行以下操作：

a = tf.zeros((5, 2), dtype=tf.int32) 
idx = tf.constant([0, 1, 1, 0, 1])
row_idx = tf.range(5)
indices = tf.transpose([row_idx, idx])
a = tf.tensor_scatter_nd_update(a, indices, row_idx)

TensorFlow使用数组索引将Tensor分配给Tensor

2 个答案: