问题

我有一个大小为`NxK`（`K≥2`）的PyTorch布尔张量，其中每行保证至少有2个`True`值。
对于每一行，我想要获取具有`True`值的2个**随机**单元格的索引。

例如，假设我有以下张量：
```python
tensor([[False,  True, False, False,  True],
        [ True,  True,  True, False,  True],
        [ True,  True, False,  True, False],
        [False,  True, False, False,  True],
        [False,  True, False, False,  True]])

因此，可能的输出是：

tensor([[1, 4],
        [4, 2],
        [0, 3],
        [4, 1],
        [4, 1]])

因为在第一行，我们在第1和第4列有True值，而在第二行，我们在第4和第2列有True值，依此类推。

另一个可能的输出是：

tensor([[4, 1],
        [1, 4],
        [0, 1],
        [1, 4],
        [1, 4]])

因为现在每行选择了不同的随机True值。

我目前为此实现了一个pytorch-numpy-pytorch解决方案，使用了np.random.choice：

available_trues = [t.nonzero(as_tuple=False).flatten() for t in input_tensor]
only_2_trues = [np.random.choice(t.cpu(), size=(2,), replace=False) for t in available_trues]
only_2_trues = torch.from_numpy(np.stack(only_2_trues)).cuda()

但是这个解决方案不是矢量化的（因为它使用了列表），而且它需要在CPU和GPU之间来回移动数据。由于我在大型矩阵上工作，并且执行此操作多次，这会导致严重减速。

我想知道是否可以在不使用列表推导或不移动数据的情况下完成这个操作（甚至两者都不用：]）。

<details>
<summary>英文:</summary>

I&#39;m given a PyTorch boolean tensor of size `NxK` (`K&gt;=2`), which is guaranteed to have at least 2 `True` values in each row.
For each row, I want to get the indices of 2 **random** cells with `True` values.

For example, let&#39;s say I have the following tensor:

tensor([[False, True, False, False, True],
[ True, True, True, False, True],
[ True, True, False, True, False],
[False, True, False, False, True],
[False, True, False, False, True]])

So a possible output would be:

tensor([[1, 4],
[4, 2],
[0, 3],
[4, 1],
[4, 1]])

Because on the first row, we have `True` value on column 1 and 4, and on the second row we have `True` value on column 4 and 2, and so on.

Another possible output would be:

tensor([[4, 1],
[1, 4],
[0, 1],
[1, 4],
[1, 4]])

Because now a different random `True` values was selected for each row.

I've currently implemented a pytorch-numpy-pytorch solution for that, using `np.random.choice`:

available_trues = [t.nonzero(as_tuple=False).flatten() for t in input_tensor]
only_2_trues = [np.random.choice(t.cpu(), size=(2,), replace=False) for t in available_trues]
only_2_trues = torch.from_numpy(np.stack(only_2_trues)).cuda()

But this solution is not vectorized (since it works using lists) and it requires moving the data back and forth between the CPU and the GPU. Since I'm working on large matrices and running this operation many times, this causes a major slowdown.

I wonder if it can be done without list comprehension or without moving the data (or even without both :] )

</details>


# 答案1
**得分**: 3

Initially, I was thinking if this can be done with `torch.topk` but then realized that it will be deterministic.

我最初考虑是否可以使用`torch.topk`来完成，但后来意识到这将是确定性的。

I was then able to make it work using `torch.multinomial` with some extra setup of a probability matrix.

然后，我能够通过使用`torch.multinomial`以及一些额外的概率矩阵设置来使其工作。

Assume the matrix is `m`.

假设矩阵为`m`。

```python
p = matrix / matrix.sum(axis=1).unsqueeze(1)

tensor([[0.0000, 0.5000, 0.0000, 0.0000, 0.5000],
        [0.2500, 0.2500, 0.2500, 0.0000, 0.2500],
        [0.3333, 0.3333, 0.0000, 0.3333, 0.0000],
        [0.0000, 0.5000, 0.0000, 0.0000, 0.5000],
        [0.0000, 0.5000, 0.0000, 0.0000, 0.5000]])

Then

然后

p.multinomial(num_samples=2)

tensor([[4, 1],
        [2, 4],
        [1, 0],
        [1, 4],
        [1, 4]])

Each time you run, you get different results.

每次运行时，您会得到不同的结果。

Obviously you can combine the above steps into one, I'm just showcasing what exact is the p matrix doing.

显然，您可以将上述步骤合并为一步，我只是展示了p矩阵的确切作用。

p = matrix / matrix.sum(axis=1).unsqueeze(1)

tensor([[0.0000, 0.5000, 0.0000, 0.0000, 0.5000],
        [0.2500, 0.2500, 0.2500, 0.0000, 0.2500],
        [0.3333, 0.3333, 0.0000, 0.3333, 0.0000],
        [0.0000, 0.5000, 0.0000, 0.0000, 0.5000],
        [0.0000, 0.5000, 0.0000, 0.0000, 0.5000]])

p.multinomial(num_samples=2)

tensor([[4, 1],
        [2, 4],
        [1, 0],
        [1, 4],
        [1, 4]])