Pytorch：有效地设置块对角矩阵？

Question

我有一个大小为 [N x 3 x 3] 的张量 A 和一个大小为 [N*3 x N*3] 的矩阵 B

我想复制A -> B的内容，这样对角线元素基本都被填满了，我想高效地做到这一点：

它应该填充 B 看起来像这样填充的东西：

所以每个 [i,3,3] 填充到 B 中的每个 [3x3] 部分沿对角线向下。

我该怎么做？尽可能高效地用于实时应用程序。我可以编写一个 CUDA 内核来执行此操作，但我更愿意使用一些特殊的 Pytorch 函数来完成它

Answer 1

使用torch.block_diag():

# Setup
A = torch.ones(3,3,3, dtype=int)

# Unpack blocks and apply
B = torch.block_diag(*A)

>>> B
tensor([[1, 1, 1, 0, 0, 0, 0, 0, 0],
        [1, 1, 1, 0, 0, 0, 0, 0, 0],
        [1, 1, 1, 0, 0, 0, 0, 0, 0],
        [0, 0, 0, 1, 1, 1, 0, 0, 0],
        [0, 0, 0, 1, 1, 1, 0, 0, 0],
        [0, 0, 0, 1, 1, 1, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 1, 1, 1],
        [0, 0, 0, 0, 0, 0, 1, 1, 1],
        [0, 0, 0, 0, 0, 0, 1, 1, 1]])

Answer 2

这是一个简单的（就地）示例，不确定真正大张量的性能：

代码：

import torch

# Create some tensors
N = 3
A = torch.ones(N, 3, 3)
A[1] *= 2
A[2] *= 3
B = torch.zeros(N*3, N*3)


def diagonalizer(A, B):
    N = A.shape[0]
    i_min = 0
    j_min = 0
    i_max = 3
    j_max = 3

    for t in range(N):
        B[i_min:i_max, j_min:j_max] = A[t]  # NOTE! this is inplace operation

        # do the step:
        i_min += 3
        j_min += 3
        i_max += 3
        j_max += 3


print('before:\n', B, sep='')

diagonalizer(A, B)

print('after:\n', B, sep='')

输出：

before:
tensor([[0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0.]])
after:
tensor([[1., 1., 1., 0., 0., 0., 0., 0., 0.],
        [1., 1., 1., 0., 0., 0., 0., 0., 0.],
        [1., 1., 1., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 2., 2., 2., 0., 0., 0.],
        [0., 0., 0., 2., 2., 2., 0., 0., 0.],
        [0., 0., 0., 2., 2., 2., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 3., 3., 3.],
        [0., 0., 0., 0., 0., 0., 3., 3., 3.],
        [0., 0., 0., 0., 0., 0., 3., 3., 3.]])