python - TensorFlow 中张量值的条件赋值

Question

我想复制以下 numpy tensorflow 中的代码.例如，我想分配一个 0到以前值为 1 的所有张量索引.

a = np.array([1, 2, 3, 1])
a[a==1] = 0

# a should be [0, 2, 3, 0]

如果我在 tensorflow 中编写类似的代码我收到以下错误。

TypeError: 'Tensor' object does not support item assignment

方括号中的条件应该是任意的，如a[a<1] = 0 .

有没有办法在 tensorflow 中实现这个“条件赋值”(因为没有更好的名字) ？

Answer 1

比较运算符，例如 greater than在 TensorFlow API 中可用。

但是，在直接操作张量方面，没有什么能比得上简洁的 NumPy 语法。您必须使用单独的 comparison、where 和 assign 运算符来执行相同的操作。

与您的 NumPy 示例等效的代码是这样的:

import tensorflow as tf

a = tf.Variable( [1,2,3,1] )    
start_op = tf.global_variables_initializer()    
comparison = tf.equal( a, tf.constant( 1 ) )    
conditional_assignment_op = a.assign( tf.where (comparison, tf.zeros_like(a), a) )

with tf.Session() as session:
    # Equivalent to: a = np.array( [1, 2, 3, 1] )
    session.run( start_op )
    print( a.eval() )    
    # Equivalent to: a[a==1] = 0
    session.run( conditional_assignment_op )
    print( a.eval() )

# Output is:
# [1 2 3 1]
# [0 2 3 0]

打印语句当然是可选的，它们只是为了证明代码正确执行。