python - 使用 tf.where() 通过 2d 条件选择 3d 张量并用键和值替换 2d 索引中的元素

Question

标题中有2个问题。我对这两个问题感到困惑，因为 tensorflow 是一种静态编程语言(我真的想回到 pytorch 或 chainer)。

我举2个例子。请用tensorflow代码回答我或提供相关函数链接。

1) tf.where()

data0 = tf.zeros([2, 3, 4], dtype = tf.float32)
data1 = tf.ones([2, 3, 4], dtype = tf.float32)
cond = tf.constant([[0, 1, 1], [1, 0, 0]])
# cond.shape == (2, 3)
# tf.where() works for 1d condition with 2d data, 
# but not for 2d indices with 3d tensor
# currently, what I am doing is:
#    cond = tf.stack([cond] * 4, 2)
data = tf.where(cond > 0, data1, data0)
# data should be [[0., 1., 1.], [1., 0., 0.]]

(我不知道如何将 cond 广播到 3d 张量)

2) 改变二维张量中的元素

# all dtype == tf.int64
t2d = tf.Variable([[0, 1, 2], [3, 4, 5]])
k, v = tf.constant([[0, 2], [1, 0]]), tf.constant([-2, -3])
# TODO: change values at positions k to v
# I cannot do [t2d.copy()[i] = j for i, j in k, v]
t3d == [[[0, 1, -2], [3, 4, 5]],
        [[0, 1, 2], [-3, 4, 5]]]

提前非常感谢您。 XD

Answer 1

这是两个完全不同的问题，它们可能应该这样发布，但无论如何。

1)

是的，您需要手动将所有输入广播到 [tf.where]( https://www.tensorflow.org/api_docs/python/tf/where] 如果它们不同。对于有值(value)的东西，有一个(旧的) open issue about it ，但到目前为止隐式广播尚未实现。您可以像您建议的那样使用 tf.stack ，尽管 tf.tile 可能会更明显(并且可能节省内存，尽管我不确定它是如何实现的):

cond = tf.tile(tf.expand_dims(cond, -1), (1, 1, 4))

或者简单地使用 tf.broadcast_to :

cond = tf.broadcast_to(tf.expand_dims(cond, -1), tf.shape(data1))

2)

这是实现此目的的一种方法:

import tensorflow as tf

t2d = tf.constant([[0, 1, 2], [3, 4, 5]])
k, v = tf.constant([[0, 2], [1, 0]]), tf.constant([-2, -3])
# Tile t2d
n = tf.shape(k)[0]
t2d_tile = tf.tile(tf.expand_dims(t2d, 0), (n, 1, 1))
# Add aditional coordinate to index
idx = tf.concat([tf.expand_dims(tf.range(n), 1), k], axis=1)
# Make updates tensor
s = tf.shape(t2d_tile)
t2d_upd = tf.scatter_nd(idx, v, s)
# Make updates mask
upd_mask = tf.scatter_nd(idx, tf.ones_like(v, dtype=tf.bool), s)
# Make final tensor
t3d = tf.where(upd_mask, t2d_upd, t2d_tile)
# Test
with tf.Session() as sess:
    print(sess.run(t3d))

输出:

[[[ 0  1 -2]
  [ 3  4  5]]

 [[ 0  1  2]
  [-3  4  5]]]