Pytorch/NumPy 批量子矩阵索引

Question

有一个形状为 (N, N)

的单一源（方形）矩阵 L

import torch as pt
import numpy as np

N = 4
L = pt.arange(N*N).reshape(N, N)  # or np.arange(N*N).reshape(N, N)
L = tensor([[ 0,  1,  2,  3],
            [ 4,  5,  6,  7],
            [ 8,  9, 10, 11],
            [12, 13, 14, 15]])

和形状为 (K, N) 的布尔掩码 m 的矩阵（向量的向量），我想根据它从 L.

中提取子矩阵

K = 3
m = tensor([[ True,  True, False, False],
            [False,  True,  True, False],
            [False,  True, False,  True]])

我知道如何通过为任何 i 调用 L[m[i]][:, m[i]] 使用单个掩码向量提取单个子矩阵。因此，例如，对于 i=0，我们会得到

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

但我需要沿着整个“批处理”维度执行操作。我正在寻找的最终结果可以通过

实现

res = []
for i in range(K):
    res.append(L[m[i]][:, m[i]])
output = pt.stack(res)

但是，我希望有一个更好的解决方案，除了for循环。我意识到如果 m 沿最后一个维度 (dim/axis=1) 的总和不是常数，for 循环解决方案本身会崩溃，但如果我能保证它是常数，是否有更好的解决方案？如果没有，更改选择器表示会有帮助吗？为了方便，我选择了布尔掩码，但我更喜欢更好的性能。

Answer 1

请注意，您可以通过索引和广播一起获取第一个方块：

r = torch.tensor([0,1])
L[r[:,None], r]

输出：

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

同样的原理可以应用到第二个方块：

r = torch.tensor([1,2])
L[r[:,None], r]

输出：

tensor([[ 5,  6],
        [ 9, 10]])

结合起来你会得到：

i = torch.tensor([[0, 1], [1, 2]])
L[i[:,:,None], i[:,None]]

输出：

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

        [[ 5,  9],
         [ 6, 10]]])

所有 3 个方块：

i = torch.tensor([
    [0, 1],
    [1, 2],
    [1, 3],
])
L[i[:,:,None], i[:,None]]

输出：

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

        [[ 5,  6],
         [ 9, 10]],

        [[ 5,  7],
         [13, 15]]])

总而言之，我建议使用索引而不是掩码。

Pytorch/NumPy 批量子矩阵索引

Pytorch/NumPy batched submatrix indexing

python

numpy

matrix

torch