如何在Python中高效计算特殊稀疏3D张量与向量的乘法？

我有一个尺寸为(nrows, ncols, ncols)的3D张量T，需要和尺寸为nrows的向量x相乘。但T是大型稀疏张量，每行仅含6×6个非零元素，所以用稀疏格式存储更高效。我定义了尺寸为(nrows, 6, 6)的张量Tsparse，以及尺寸为(nrows, 6)的指针矩阵p来存储每行非零元素的索引。现在需要用Python代码实现Tsparse与x的乘法，优先使用NumPy或SciPy实现。以下是我的最小可行示例（MWE），需要基于需求改写：

import numpy as np

nrows = 100
ncols = 20
p = np.zeros((nrows, 6), dtype=int)

for i in range(nrows):
    p[i, :] = np.random.choice(np.arange(ncols), 6, replace=False)

T = np.zeros((nrows, ncols, ncols))
Tsparse = np.zeros((nrows, 6, 6))
for i in range(nrows):
    for j in range(6):
        for k in range(6):
            entry = np.random.rand()
            T[i, p[i, j], p[i, k]] = entry
            Tsparse[i, j, k] = entry
            
x = np.array(np.random.rand(nrows))

res = np.tensordot(T, x, axes=([0], [0]))

#The goal would be to form Tsparse * x instead of T * x

print(res)

直接用NumPy就能高效实现该计算，无需构建完整的稠密张量T，核心是通过广播和索引操作，将Tsparse与x的加权和映射到最终结果矩阵中。

基础实现版本 #

import numpy as np

nrows = 100
ncols = 20
p = np.zeros((nrows, 6), dtype=int)

# 生成随机不重复列索引
for i in range(nrows):
    p[i, :] = np.random.choice(np.arange(ncols), 6, replace=False)

# 生成稀疏张量块
Tsparse = np.zeros((nrows, 6, 6))
for i in range(nrows):
    for j in range(6):
        for k in range(6):
            Tsparse[i, j, k] = np.random.rand()
            
x = np.random.rand(nrows)

# 核心计算：对每个稀疏块加权，再累加到对应位置
weighted_blocks = Tsparse * x[:, np.newaxis, np.newaxis]
res_sparse = np.zeros((ncols, ncols), dtype=np.float64)

for i in range(nrows):
    idx = p[i]
    res_sparse[np.ix_(idx, idx)] += weighted_blocks[i]

# 验证结果与稠密计算一致
T = np.zeros((nrows, ncols, ncols))
for i in range(nrows):
    idx = p[i]
    T[i, np.ix_(idx, idx)] = Tsparse[i]
res_dense = np.tensordot(T, x, axes=([0], [0]))
print("结果最大误差：", np.max(np.abs(res_sparse - res_dense)))

向量化优化版本 #

如果nrows极大，可替换循环为NumPy向量化操作，进一步提升效率：

import numpy as np

nrows = 100
ncols = 20
p = np.zeros((nrows, 6), dtype=int)

for i in range(nrows):
    p[i, :] = np.random.choice(np.arange(ncols), 6, replace=False)

Tsparse = np.random.rand(nrows, 6, 6)
x = np.random.rand(nrows)

# 加权处理
weighted_blocks = Tsparse * x[:, np.newaxis, np.newaxis]

# 展开所有非零元素的坐标和值
row_indices = p[:, np.newaxis, :].repeat(6, axis=1).flatten()
col_indices = p[:, :, np.newaxis].repeat(6, axis=2).flatten()
values = weighted_blocks.flatten()

# 批量累加
res_sparse_vec = np.zeros((ncols, ncols), dtype=np.float64)
np.add.at(res_sparse_vec, (row_indices, col_indices), values)

# 验证
T = np.zeros((nrows, ncols, ncols))
for i in range(nrows):
    idx = p[i]
    T[i, np.ix_(idx, idx)] = Tsparse[i]
res_dense = np.tensordot(T, x, axes=([0], [0]))
print("向量化版本最大误差：", np.max(np.abs(res_sparse_vec - res_dense)))

优势说明 #

• 内存占用从稠密张量的O(nrows*ncols²)降至O(nrows*6²)，适合ncols远大于6的场景。
• 两种实现都完全基于NumPy，无需额外依赖，计算效率远高于先构建稠密张量再运算的方式。

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