来自多维张量的加权采样

Question

我需要对多维张量进行加权采样。

我有一个形状为[X,Y] 的张量A 和一个形状为[X] 的概率分布B。 我需要根据分布B 从A 中采样N 元素。

B 表示子张量的分布。每个子张量内的采样是均匀的。

A 中有一些填充，所以我必须考虑到这一点。什么是填充的信息包含在掩码中。

例如

A      = [[1,   2,   3,   X,   X,  X],
          [10,  20,  30,  X,   X,  X],
          [100, 200, 300, 400, 500, 600]]
A_mask = [[T,   T,   T,   F,   F,  F],
          [T,   T,   T,   F,   F,  F],
          [T,   T,   T,   T,   T,  T]]
B = [0.5, 0.4, 0.1]

# a possible output, with N = 10
ouput = [1, 1, 2, 2, 3, 10, 20, 30, 30, 200]

我能够从A 的每个嵌套张量中检索要采样的元素数量：

tf.multinomial(tf.log(probability_distribution), N)

# a possible output of that function, with N = 10, is:
[1, 1, 1, 1, 1, 2, 2, 2, 2, 3]

对于这些数字中的每一个，我必须在该子张量中执行均匀采样。

我能够计算每个子张量的 最大值。

subtensor_sizes = tf.reduce_sum(tf.cast(A_mask, tf.int32), axis=1)

# it would return: [3, 3, 6]

此时，对于多项式函数返回的每个子张量，我应该在0 和它的maxvalue 之间执行统一采样（或者类似地，计算出现次数并从出现@987654336 的子张量中采样T 元素@times 在多项式的输出中）。

我不知道该怎么做，怎么办？

Answer 1

所以你有一个张量A 包含不同长度的序列。您想从这些序列中提取值，以不同的概率B 为每个序列选择一个值。

您可以按照以下方式进行：

import tensorflow as tf

A = tf.constant(
    [[1,   2,   3,   -1,  -1,  -1],
     [10,  20,  30,  -1,  -1,  -1],
     [100, 200, 300, 400, 500, 600]])
A_mask = tf.constant(
    [[True,   True,   True,   False,   False,  False],
     [True,   True,   True,   False,   False,  False],
     [True,   True,   True,   True,   True,  True]])
B = tf.constant([0.5, 0.4, 0.1])
subtensor_sizes = tf.reduce_sum(tf.cast(A_mask, tf.int32), axis=1)

# get random sample index
output = tf.to_int32(tf.multinomial(tf.log(B[None]), 10)[0])
# get corresponding sample size
output_sizes = tf.gather(subtensor_sizes, output)
# generate a random index in each range
random_idxs = tf.map_fn(
  lambda x: tf.random_uniform((), maxval=x, dtype=tf.int32), output_sizes)
# construct nd-index for tf.gather
random_ndxs = tf.concat([output[:, None], random_idxs[:, None]], axis=-1)
# get sample values
random_samples = tf.gather_nd(A, random_ndxs)