为何Numpy中3D数组切片的矩阵乘法速度极慢？

问题 #

我在使用Numpy进行大维度矩阵乘法时遇到性能差异：生成3D数组Y1，Y2是Y1[:,:,0]的副本，Y3是独立生成的2D数组，与矩阵W相乘时，Y1[:,:,0].dot(W)耗时约34秒，而Y2.dot(W)、Y3.dot(W)仅约0.06秒。猜测是3D数组切片的性能问题，且生成Y2的副本也较慢。现需处理给定的3D数组，反复计算其切片与不同W的矩阵乘积，求更高效的实现方式。

测试代码：

import numpy as np
from time import time

Y1 = np.random.uniform(-1, 1, (5000, 1093, 201))
Y2 = Y1[:,:,0].copy()
Y3 = np.random.uniform(-1, 1, (5000, 1093))

W = np.random.uniform(-1, 1, (1093, 30))

# method 1
START = time()
Y1[:,:,0].dot(W)
END = time()
print(f"Method 1 : {END - START}")

# method 2
START = time()
Y2.dot(W)
END = time()
print(f"Method 2 : {END - START}")

# method 3
START = time()
Y3.dot(W)
END = time()
print(f"Method 3 : {END - START}")

测试结果：

• Method 1：约34秒
• Method 2：约0.06秒
• Method 3：约0.06秒

高效实现方案 #

核心原因 #

Y1[:,:,0]返回的是数组视图（view），而非独立副本。由于3D数组默认按C顺序存储（最后一维优先连续），切片后的视图内存是不连续的，而Numpy的dot运算对连续内存数组有显著优化，这是性能差异的根源。

方法1：将切片转为连续内存数组（单一切片反复使用场景） #

使用np.ascontiguousarray将切片视图转为连续内存数组，无需完整复制（若原视图已连续则直接返回），后续反复用该数组与不同W相乘即可：

# 预处理切片，仅需执行一次
Y_slice = np.ascontiguousarray(Y1[:, :, 0])

# 后续反复计算（示例）
W_list = [np.random.uniform(-1,1,(1093,30)) for _ in range(10)]
for W in W_list:
    result = Y_slice.dot(W)
    # 处理结果

该方法预处理耗时远低于直接用视图计算，后续每次乘法均能达到Y2.dot(W)的性能。

方法2：批量矩阵乘法（多切片处理场景） #

若需处理Y1的多个切片（如所有201层），利用Numpy的广播和批量矩阵乘法优化，一次性完成所有计算：

# 将Y1的后两维转置，变为(5000, 201, 1093)，使每一层切片内存连续
Y1_transposed = Y1.transpose(0, 2, 1)
# 批量乘法，结果为(5000, 201, 30)，对应每一层切片与W的乘积
batch_result = Y1_transposed @ W

# 若仅需第0层结果，直接索引
result_0 = batch_result[:, 0, :]

此方法避免了循环切片，利用Numpy底层优化大幅提升整体效率，尤其适合多切片批量处理场景。

方法3：先全局运算再取切片（单一切片场景备选） #

直接对整个3D数组Y1与W做矩阵乘法，再提取目标切片，利用全局运算的优化特性：

# 全局运算，结果为(5000, 201, 30)
full_result = np.dot(Y1, W).transpose(0, 2, 1)
# 提取第0层切片的结果
result_0 = full_result[:, 0, :]

该方法无需预处理，适合偶尔计算单一切片的场景，性能接近连续数组的乘法效率。

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