Tensorflow，tf.multinomial，获取关联概率失败

Question

我正在尝试使用tf.multinomial进行采样，我希望获得采样值的相关概率值。下面是我的示例代码，

​

In [1]: import tensorflow as tf

In [2]: tf.enable_eager_execution()

In [3]: probs = tf.constant([[0.5, 0.2, 0.1, 0.2], [0.6, 0.1, 0.1, 0.1]], dtype=tf.float32)

In [4]: idx = tf.multinomial(probs, 1)

In [5]: idx  # print the indices
Out[5]:
<tf.Tensor: id=43, shape=(2, 1), dtype=int64, numpy=
array([[3],
       [2]], dtype=int64)>

In [6]: probs[tf.range(probs.get_shape()[0], tf.squeeze(idx)]
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-8-56ef51f84ca2> in <module>
----> 1 probs[tf.range(probs.get_shape()[0]), tf.squeeze(idx)]

C:\ProgramData\Anaconda3\lib\site-packages\tensorflow\python\ops\array_ops.py in _slice_helper(tensor, slice_spec, var)
    616       new_axis_mask |= (1 << index)
    617     else:
--> 618       _check_index(s)
    619       begin.append(s)
    620       end.append(s + 1)

C:\ProgramData\Anaconda3\lib\site-packages\tensorflow\python\ops\array_ops.py in _check_index(idx)
    514     # TODO(slebedev): IndexError seems more appropriate here, but it
    515     # will break `_slice_helper` contract.
--> 516     raise TypeError(_SLICE_TYPE_ERROR + ", got {!r}".format(idx))
    517
    518

TypeError: Only integers, slices (`:`), ellipsis (`...`), tf.newaxis (`None`) and scalar tf.int32/tf.int64 tensors are valid indices, got <tf.Tensor: id=7, shape=(2,), dtype=int32, numpy=array([3, 2])>

如idx所示，我想要的预期结果是[0.2, 0.1]。

但在Numpy中，此方法的工作方式与https://stackoverflow.com/a/23435869/5046896中的回答相同

我怎么才能修复它？

Answer 1

你可以试试tf.gather_nd，你可以试试

>>> import tensorflow as tf
>>> tf.enable_eager_execution()
>>> probs = tf.constant([[0.5, 0.2, 0.1, 0.2], [0.6, 0.1, 0.1, 0.1]], dtype=tf.float32)
>>> idx = tf.multinomial(probs, 1)
>>> row_indices = tf.range(probs.get_shape()[0], dtype=tf.int64)
>>> full_indices = tf.stack([row_indices, tf.squeeze(idx)], axis=1)
>>> rs = tf.gather_nd(probs, full_indices)

或者，您可以使用tf.distributions.Multinomial，优点是您不需要关心上面代码中的batch_size。当您设置batch_size=None时，它可以在不同的batch_size下工作。这里有一个简单的例子，

multinomail = tf.distributions.Multinomial(
                         total_count=tf.constant(1, dtype=tf.float32),  # sample one for each record in the batch, that is [1, batch_size]
                         probs=probs)
sampled_actions = multinomail.sample()  # sample one action for data in the batch
predicted_actions = tf.argmax(sampled_actions, axis=-1) 
action_probs = sampled_actions * predicted_probs 
action_probs = tf.reduce_sum(action_probs, axis=-1)

我更喜欢后者，因为它灵活而优雅。