如何在PyTorch中无循环实现基于累积乘积的特定矩阵生成

可以用PyTorch的向量化操作完全替代循环，利用前缀累积乘积和广播机制实现，性能远优于Python循环，适合在模型forward()中使用。

无循环实现代码 #

import torch
import matplotlib.pyplot as plt

def gen_matrix_no_loop(array):
    N = array.shape[0] + 1
    # 生成包含初始值1的前缀累积乘积数组
    prefix = torch.cat([
        torch.tensor([1.0], device=array.device, dtype=array.dtype),
        torch.cumprod(array, dim=0)
    ])
    # 利用广播计算上三角矩阵（含对角线）
    upper_tri = torch.triu(prefix[None, :] / prefix[:, None], diagonal=0)
    # 构造对称矩阵：上三角 + 转置上三角 - 对角线（避免重复计算）
    mat = upper_tri + upper_tri.T - torch.diag(upper_tri.diag())
    return mat

# 验证与原循环实现的一致性
example_array = torch.tensor([0.1, 0.2, 0.05, 0.3, 0.2])
output_matrix_loop = gen_matrix(example_array)
output_matrix_no_loop = gen_matrix_no_loop(example_array)

print("结果是否一致：", torch.allclose(output_matrix_loop, output_matrix_no_loop))  # 输出 True

# 可视化
plt.figure()
plt.imshow(output_matrix_no_loop)
plt.colorbar()
plt.show()

实现原理 #

1. 前缀累积乘积：
   构造prefix数组，其中prefix[0] = 1，prefix[k]表示array[0]到array[k-1]的累积乘积。这样任意区间array[i:j]的累积乘积可以转化为prefix[j] / prefix[i]，避免逐行计算累积乘积。

2. 广播与三角矩阵：
   通过prefix[None, :] / prefix[:, None]利用广播生成(N×N)矩阵，其中每个元素(i,j)的值正好对应原循环中m[i,j]（当j >= i时）。再用torch.triu保留上三角部分（含对角线）。

3. 对称矩阵构造：
   由于目标矩阵是对称的，将上三角矩阵转置后与原上三角矩阵相加，再减去对角线（避免对角线被重复计算两次），最终得到完整的对称矩阵。

优势 #

• 完全基于PyTorch内置的向量化操作，无Python循环，GPU加速下性能大幅提升；
• 代码简洁，可直接嵌入模型的forward()函数中，避免循环带来的性能瓶颈。

内容的提问来源于stack exchange，提问作者Carciofo