python - 在 PyTorch 中并行求和特征图中的所有对角线

Question

假设我有一个形状为 (1, 64, 128, 128) 的张量，并且我想创建一个形状为 (1, 64, 255) 的张量每个 (128, 128) 矩阵的所有对角线之和(有 1 个主矩阵，下面有 127 个对角线，上面有 127 个对角线，所以总共 255 个)。我目前正在做的事情如下:

x = torch.rand(1, 64, 128, 128)

diag_sums = torch.zeros(1, 64, 255)
j = 0
for k in range(-127, 128):
    diag_sums[j, :, k + 127] = torch.diagonal(x, offset=k, dim1=-2, dim2=-1).sum(dim=2)

这显然非常慢，因为它使用 Python 循环并且不是相对于 k 并行完成的。

我认为这不能使用 torch.diagonal 来完成，因为该函数明确使用单个 int 作为偏移参数。如果我可以在那里传递一个列表，这会起作用，但我想实现起来会很复杂(需要 PyTorch 本身的更改)。

我认为可以使用torch.einsum来实现这一点，但我想不出一种方法来做到这一点。

这就是我的问题:如何获得上述张量？

Answer 1

您是否考虑过使用torch.nn.functional.conv2d ？
您可以使用在张量上滑动的对角过滤器和适当的零填充来对对角线求和。

import torch
import torch.nn.functional as nnf

# construct a diagonal filter using `eye` function, shape it appropriately
f = torch.eye(x.shape[2])[None, None,...].repeat(x.shape[1], 1, 1, 1)
# compute the diagonal sum with appropriate zero padding
conv_diag_sums = nnf.conv2d(x, f, padding=(x.shape[2]-1,0), groups=x.shape[1])[..., 0]

请注意，结果的顺序与您在循环中计算的顺序略有不同:

diag_sums = torch.zeros(1, 64, 255)
for k in range(-127, 128):
    diag_sums[j, :, 127-k] = torch.diagonal(x, offset=k, dim1=-2, dim2=-1).sum(dim=2)

# compare
(conv_diag_sums == diag_sums).all()

结果为True - 它们是相同的。