PyTorch tensor.sum()张量超特定尺寸性能骤降，原因何在？

问题复现 #

测试显示，当张量宽度从1000扩展到2000时，sum(dim=0)操作耗时从0.69s飙升至20s，而NumPy执行相同操作仅需0.4s：

PyTorch测试代码 #

import torch
import time

# 生成200000x2000的Bool张量
tensor = torch.FloatTensor(200_000, 2_000).uniform_() > 0.8

# 测试1000宽度的sum
start = time.time()
tensor[:, :1000].sum(dim=0)
print(f"1000宽度耗时: {time.time() - start:.2f}s")  # 输出0.69s

# 测试2000宽度的sum
start = time.time()
tensor[:, :2000].sum(dim=0)
print(f"2000宽度耗时: {time.time() - start:.2f}s")  # 输出20s

NumPy对比测试 #

import numpy as np
import time

start = time.time()
np.array(tensor[:, :2000]).sum(axis=0)
print(f"NumPy 2000宽度耗时: {time.time() - start:.2f}s")  # 输出0.40s

核心原因分析 #

1. Bool张量的优化差异 #

PyTorch对CPU上Bool类型张量的sum操作优化远不如NumPy。Bool张量的位存储特性导致求和时需要额外的类型转换或位运算逻辑，而NumPy的底层BLAS实现针对这类操作做了更高效的硬件适配，能快速完成布尔值的累加。

2. CPU缓存命中率下降 #

当张量宽度从1000增加到2000时，单维度的元素数量翻倍，超出了CPU L2/L3缓存的容纳范围。PyTorch的内存访问模式在这种情况下会频繁触发缓存 miss，导致大量时间浪费在内存读写上；而NumPy的内存布局和访问策略更贴合CPU缓存机制，缓存命中率更高。

3. 自动微分的额外开销 #

默认情况下，PyTorch张量会追踪计算图（即使没有反向传播需求），这会为sum操作带来额外的簿记开销。NumPy没有计算图追踪机制，操作更轻量化，尤其在大张量场景下差异明显。

4. 并行化策略不同 #

NumPy的sum操作会充分利用CPU多核心并行计算，且针对大张量的并行粒度做了优化；而PyTorch对Bool张量的sum可能未启用最优的并行策略，导致大张量下无法有效利用多核资源。

验证与优化方案 #

• 转换张量类型：将Bool张量转为Float张量后再求和，能大幅提升性能：

  tensor.float()[:, :2000].sum(dim=0)
  

• 关闭计算图追踪：用torch.no_grad()包裹操作，消除自动微分的额外开销：

  with torch.no_grad():
      tensor[:, :2000].sum(dim=0)
  

• 确保MKL后端启用：安装带MKL优化的PyTorch版本（官方默认安装包已包含），能提升CPU张量操作的性能。

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