Question

我正在尝试将切片更新的NumPy操作转换为TensorFlow。我想重现以下最小示例：

input = np.arange(3 * 5).reshape((3, 5))

array([[ 0,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14]])

input[:, [0, 2]] = -1

array([[-1,  1, -1,  3,  4],
       [-1,  6, -1,  8,  9],
       [-1, 11, -1, 13, 14]])

因此，我想为数组中某些列的所有元素设置一个常量值。

现在，我有了Tensors而不是NumPy数组，列索引也被动态计算并存储在Tensors中。我发现了如何使用tf.scatter_nd_update更新给定行中的所有值：

input = tf.Variable(tf.reshape(tf.range(3 * 5, dtype=tf.int32), [3, 5]))                                                                                                                                                                                          
indices = tf.constant([[0], [2]])                                                                                                                                                                                                                                 
updates = tf.constant([[-1, -1, -1, -1, -1], [-1, -1, -1, -1, -1]])                                                                                                                                                                                               

scatter = tf.scatter_nd_update(input, indices, updates)                                                                                                                                                                                                           

with tf.Session() as sess:                                                                                                                                                                                                                                        
    sess.run(tf.global_variables_initializer())                                                                                                                                                                                                                   
    print(sess.run(scatter))

输出：

[[-1 -1 -1 -1 -1]
 [ 5  6  7  8  9]
 [-1 -1 -1 -1 -1]]

但是如何对某些列执行此操作？

Answer 1

您可以这样做：

import tensorflow as tf

def update_columns(variable, columns, value):
    columns = tf.convert_to_tensor(columns)
    rows = tf.range(tf.shape(variable)[0], dtype=columns.dtype)
    ii, jj = tf.meshgrid(rows, columns, indexing='ij')
    value = tf.broadcast_to(value, tf.shape(ii))
    return tf.scatter_nd_update(variable, tf.stack([ii, jj], axis=-1), value)

inp = tf.Variable(tf.reshape(tf.range(3 * 5, dtype=tf.int32), [3, 5]))
updated = update_columns(inp, [0, 2], -1)
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    print(sess.run(updated))

输出：

[[-1  1 -1  3  4]
 [-1  6 -1  8  9]
 [-1 11 -1 13 14]]

但是请注意，如果您确实想使用变量（并为其分配新值），则仅应使用tf.scatter_nd_update。如果要获得与另一个张量相等但已更新某些值的张量，则应使用常规的张量操作，而不要将其转换为变量。例如，对于这种情况，您可以执行以下操作：

import tensorflow as tf

def update_columns_tensor(tensor, columns, value):
    columns = tf.convert_to_tensor(columns)
    shape = tf.shape(tensor)
    num_rows, num_columns = shape[0], shape[1]
    mask = tf.equal(tf.range(num_columns, dtype=columns.dtype), tf.expand_dims(columns, 1))
    mask = tf.tile(tf.expand_dims(tf.reduce_any(mask, axis=0), 0), (num_rows, 1))
    value = tf.broadcast_to(value, shape)
    return tf.where(mask, value, tensor)

inp = tf.reshape(tf.range(3 * 5, dtype=tf.int32), [3, 5])
updated = update_columns_tensor(inp, [0, 2], -1)
with tf.Session() as sess:
    print(sess.run(updated))
    # Same output

了解tf.scatter_nd_update：如何更新列值？

1 个答案: