时间戳

首页 > 解决方案 > 如何将此 Python for 循环转换为 Tensorflow（如果可能）

python - 如何将此 Python for 循环转换为 Tensorflow（如果可能）

问题描述

我想在不破坏 TF 计算图的情况下将这个非常简单的 Python for 循环转换为 Tensorflow。有什么办法可以做到吗？（如有）

这是 for 循环，其中 A、B 和 C 是张量。

  for a, b in zip(A, B):
    if 0 <= a:
      C[a] += b

我对 Tensorflow 很陌生，我知道这可能是一个不合理的问题（根据 TF 的工作原理），如果我也离现实太远，我将不胜感激。谢谢，

标签： pythonfor-looptensorflowzip

解决方案

您不需要在 tensorflow 中使用循环。tf.scatter_nd是适合您问题的功能。一个例子：

import tensorflow as tf

A = [-1,2,3,1]
B = [-2,4,6,2]
C = [-1,2,3,1,7,6]

for a, b in zip(A, B):
    if 0 <= a:
        C[a] += b
print('your result:\n',C)

A = tf.constant([-1,2,3,1])
B = tf.constant([-2,4,6,2])
C = tf.constant([-1,2,3,1,7,6])

mask_B = tf.boolean_mask(B,tf.greater_equal(A,0))
mask_A = tf.boolean_mask(A,tf.greater_equal(A,0))
C = tf.scatter_nd(tf.expand_dims(mask_A,axis=-1), mask_B, tf.shape(C))+C
# if C type is tf.Variable, you can use tf.scatter_nd_add
# C = tf.scatter_nd_add(C, tf.expand_dims(mask_A,axis=-1), mask_B)

with tf.Session() as sess:
    print('tensorflow version:\n',sess.run(C))

your result:
 [-1, 4, 7, 7, 7, 6]
tensorflow version:
 [-1  4  7  7  7  6]

推荐阅读