如何为张量流中每个值的给定概率采样张量值？

Question

We need to sample values from one tensor regarding to another tensor that contains probabilities. Lets say, we have two tensors t1,t2 of shape (?,3), and want to find another tensor t3 of shape (?,1) that contains a sample of each row in t1 regarding to probabilities in t2.

Answer 1

You do this in two steps. First, you sample indices 0 , 1 and 2 , then you replace those indices with tensor values. This can be done with tf.random.categorical (see this question for more information about this function). Note tf.random.categorical was added in version 1.13.1.

import tensorflow as tf

with tf.Graph().as_default(), tf.Session() as sess:
    tf.random.set_random_seed(0)
    values = tf.placeholder(tf.float32, [None, 3])
    probabilities = tf.placeholder(tf.float32, [None, 3])
    # You can change the number of samples per row (or make it a placeholder)
    num_samples = 1
    # Use log to get log-probabilities or give logit values (pre-softmax) directly
    logits = tf.log(probabilities)
    # Sample column indices
    col_idx = tf.random.categorical(logits, num_samples, dtype=tf.int32)
    # Make row indices
    num_rows = tf.shape(values)[0]
    row_idx = tf.tile(tf.expand_dims(tf.range(num_rows), 1), (1, num_samples))
    # Gather actual values
    result = tf.gather_nd(values, tf.stack([row_idx, col_idx], axis=-1))
    # Test
    print(sess.run(result, feed_dict={
        values: [[10., 20., 30.], [40., 50., 60.]],
        # Can be actual or "proportional" probabilities not summing 1
        probabilities: [[.1, .3, .6], [0., 4., 2.]]
    }))
    # [[30.]
    #  [60.]]