在PyTorch中通过累积乘积生成特定矩阵，无需for循环

Question

在 pytorch 中，我想从 (N-1) 元素数组开始生成一个特定的 NxN 矩阵，遵循沿着数组元素的累积乘积的特定顺序来填充矩阵的每一行和每一列。

从数组开始： 这应该是输出矩阵：

我想在没有Python for循环的情况下执行此操作，以避免模型的forward()函数显着变慢

这是一个使用 for 循环的工作示例：

import torch
import matplotlib.pyplot as plt

def gen_matrix(array):
    N = array.shape[0] + 1
    m = torch.ones(N,N)
    for i in range(N):
        m[i   ,i+1:] = torch.cumprod(array[i:],0)
        m[i+1:,i   ] = torch.cumprod(array[i:],0)
    return m

example_array = torch.tensor([0.1, 0.2, 0.05, 0.3, 0.2])
output_matrix = gen_matrix(example_array)

plt.figure()
plt.imshow(output_matrix)
plt.colorbar()

是否可以用一些pytorch方法替代循环？

我试图寻找一种合适的方法，但是我对 pytorch 仍然很陌生，无法找到一种方法来实现这一点，而不需要对每一行进行切片和循环。输出矩阵是对称的，因为它等于其转置矩阵，因此我可以使用

.T

和

.tril()

的组合生成其中的一半，但最终我仍然会在行或列上循环。

Answer 1

非常有趣，我花了一段时间才找到解决方案！
一个可能的改进是将累积乘积矢量化。

首先在前面连接一个

1

并逐行展开，这将是输出张量：

>>> m = torch.cat([torch.ones(1),array])[None].repeat(N,1)
tensor([[1.0000, 0.1000, 0.2000, 0.0500, 0.3000, 0.2000],
        [1.0000, 0.1000, 0.2000, 0.0500, 0.3000, 0.2000],
        [1.0000, 0.1000, 0.2000, 0.0500, 0.3000, 0.2000],
        [1.0000, 0.1000, 0.2000, 0.0500, 0.3000, 0.2000],
        [1.0000, 0.1000, 0.2000, 0.0500, 0.3000, 0.2000],
        [1.0000, 0.1000, 0.2000, 0.0500, 0.3000, 0.2000]])

要沿torch.cumprod

应用

dim=1，我们需要将下三角形设置为

1

：

>>> m = m.triu(1) + torch.ones_like(m).tril(0)
tensor([[1.0000, 0.1000, 0.2000, 0.0500, 0.3000, 0.2000],
        [1.0000, 1.0000, 0.2000, 0.0500, 0.3000, 0.2000],
        [1.0000, 1.0000, 1.0000, 0.0500, 0.3000, 0.2000],
        [1.0000, 1.0000, 1.0000, 1.0000, 0.3000, 0.2000],
        [1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 0.2000],
        [1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000]])

要获得所需的结果，请复制上三角形，转置它，然后添加到

m

。

总而言之，归结为：

def gen_matrix(array):
    N = len(array) + 1
    m = torch.cat([torch.ones(1),array])[None].repeat(N,1)
    m = m.triu(1) + torch.ones_like(m).tril(0)
    m = m.cumprod(1)
    return m.triu() + m.triu(1).T