从轴上的Tensorflow中的张量采样

Question

我有一个L形状矩阵(2,5,2).沿最后一个轴的值​​形成概率分布.我想采样另一个S形状矩阵,(2, 5)其中每个条目是以下整数之一:0, 1.例如,

L = [[[0.1, 0.9],[0.2, 0.8],[0.3, 0.7],[0.5, 0.5],[0.6, 0.4]],
     [[0.5, 0.5],[0.9, 0.1],[0.7, 0.3],[0.9, 0.1],[0.1, 0.9]]]

其中一个样本可能是,

S = [[1, 1, 1, 0, 1],
     [1, 1, 1, 0, 1]]

在上面的例子中,分布是二项式的.但是,通常,最后一个维度L可以是任何正整数,因此分布可以是多项式的.

需要在Tensorflow计算图中有效地生成样本.我知道如何使用numpy使用函数apply_along_axis和numpy.random.multinomial.

Answer 1

你可以tf.multinomial()在这里使用.

您首先需要将输入张量重塑为形状[-1, N](N最后一个维度L):

# L has shape [2, 5, 2]
L = tf.constant([[[0.1, 0.9],[0.2, 0.8],[0.3, 0.7],[0.5, 0.5],[0.6, 0.4]],
                 [[0.5, 0.5],[0.9, 0.1],[0.7, 0.3],[0.9, 0.1],[0.1, 0.9]]])

dims = L.get_shape().as_list()
N = dims[-1]  # here N = 2

logits = tf.reshape(L, [-1, N])  # shape [10, 2]

现在我们可以将该函数tf.multinomial()应用于logits:

samples = tf.multinomial(logits, 1)
# We reshape to match the initial shape minus the last dimension
res = tf.reshape(samples, dims[:-1])