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

时间:2018-12-05 12:09:14

标签: python tensorflow numpy-broadcasting

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

我举两个例子。请以张量流代码或提供相关功能链接回答我。

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)更改2d张量中的元素

# 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

1 个答案:

答案 0 :(得分:2)

这是两个截然不同的问题,应该以同样的方式发布,但是无论如何。

1)

是的,您需要将所有输入手动广播到[tf.wherehttps://www.tensorflow.org/api_docs/python/tf/where],如果它们不同,那么对于价值而言,有一个(旧的)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]]]
相关问题