python - 使用随机生成的整数作为tensorflow map_fn中的张量索引

Question

我在 Keras 的自定义层中初始化张量 A ，其中 batchSize 是占位符:

A = K.zeros([batchSize, 2, 2 ,2])

我还初始化了一个大小为 [3,2,2,2] 的 numpy 数组 B。我想从 B 中随机选择 [i,2,2,2] 大小的数组，其中 i = 0,1,2，并且将其分配给 A 的第一个维度并重复此batchSize 次数。

由于我无法显式循环batchSize，因此我尝试了tensorflow.map_fn，如下所示:

ANew = tf.map_fn(lambda x: K.variable(B[np.random.randint(0,3,size=(1)).tolist()[0],:,:,:],
                 A, dtype=’float’, back_prop=False, infer_shape=True)

这会生成ANew张量。然而，看起来 np.random.randint 只被调用一次；结果，我总是选择相同的索引。如何修改代码以便 np.random.randint(0,3,size=(1)).tolist()[0] 被调用batchSize 次数？

Answer 1

您正在寻找K.gather。

A = K.gather(B, indices_list)