TensorFlow 通过数组索引将张量分配给张量答案

【问题标题】：TensorFlow assign Tensor to Tensor with array indexingTensorFlow 通过数组索引将张量分配给张量
【发布时间】：2019-05-07 01:17:57
【问题描述】：

我想在 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

意思是用另一个张量索引一个二维张量的所有行，然后为其分配一个张量。我对如何实现这一目标一无所知。

到目前为止，我在 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 的功能先进，但这应该仍然可以通过某种方式实现。

【问题讨论】：

标签： python tensorflow

【解决方案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]]

【讨论】：

【解决方案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)

【讨论】：