在Tensorflow中从另一个中选择随机张量

时间:2017-10-13 02:11:37

标签: tensorflow

我有一个Tensor X形状[B,L,E](让我们说,B批L长度为E的向量)。从这个Tensor X,我想在每批中随机选择N个向量,因此创建具有[B,N,E]形状的Y.

我试图将tf.random_uniform和tf.gather结合起来,但我真的很难维度,并且不能得到Y.

1 个答案:

答案 0 :(得分:4)

您可以使用以下内容:

import tensorflow as tf
import numpy as np

B = 3
L = 5
E = 2
N = 3

input = np.array(range(B * L * E)).reshape([B, L, E])
print(input)
print("#################################")

X = tf.constant(input)
batch_range = tf.tile(tf.reshape(tf.range(B, dtype=tf.int32), shape=[B, 1, 1]), [1, N, 1])
random = tf.random_uniform([B, N, 1], minval = 0, maxval = L - 1, dtype = tf.int32)

indices = tf.concat([batch_range, random], axis = 2)

output = tf.gather_nd(X, indices)
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    print(sess.run(indices))
    print("#################################")
    print(sess.run(output))