为何numpy.add.at比简单分组求和的性能表现更差？

背景与需求 #

我正在优化自己实现的KMeans算法，手头有两个numpy数组：

• 数据集X：形状(n_samples, n_features)，数据类型np.float64
• 标签y：形状(n_samples,)，数据类型np.int64，所有元素满足0 <= y < n_clusters

我的需求是按标签分组数据集，计算每组的求和值（后续再除以数量得到聚类中心）。

简易实现（性能反而更好） #

先写了个直观的实现：

new_centers = np.array([X[y == i].sum(0) for i in range(n_clusters)])

这个实现的逻辑是循环每个簇，生成掩码筛选对应样本，求和后堆叠成结果。缺点是要循环n_clusters次，每次生成掩码、复制数据，但实际测试性能却不错。

尝试用numpy.add.at优化（结果更慢） #

我以为用numpy.add.at这种纯C实现、无Python循环的方法会更快，代码如下：

new_centers = np.zeros((n_clusters, n_features))
np.add.at(new_centers, y, X)

但在n_clusters=100, n_samples=100000, n_features=10的测试条件下，它的性能反而不如简易实现：

# 测试环境初始化
n_clusters, n_samples, n_features = 100, 100_000, 10
rng = np.random.default_rng()
X = rng.random((n_samples, n_features))
y = rng.integers(0, n_clusters, n_samples)

# 简易实现耗时
%timeit np.array([X[y == i].sum(0) for i in range(n_clusters)])
# 输出：29.5 ms ± 852 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

# numpy.add.at实现耗时
%timeit np.add.at(np.zeros((n_clusters, n_features)), y, X)
# 输出：58.7 ms ± 708 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

PyTorch对比：性能碾压numpy.add.at #

我用PyTorch实现了两种类似逻辑的方法，性能远优于numpy版本：

# 方法1：index_put_累加
%%timeit
new_centers = np.zeros((n_clusters, n_features))
torch.from_numpy(new_centers).index_put_([torch.from_numpy(y)],
                                         torch.from_numpy(X), accumulate=True)
# 输出：1.74 ms ± 3.93 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

# 方法2：index_add_
%%timeit
new_centers = np.zeros((n_clusters, n_features))
torch.from_numpy(new_centers).index_add_(0, torch.from_numpy(y), torch.from_numpy(X))
# 输出：8.97 ms ± 51.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

版本信息 #

sys.version, np.__version__, torch.__version__
# 输出：
('3.11.3 (tags/v3.11.3:f3909b8, Apr  4 2023, 23:49:59) [MSC v.1934 64 bit (AMD64)]',
 '1.24.3',
 '2.0.0+cu118')

为什么numpy.add.at性能这么差？ #

1. 随机内存访问导致缓存命中率极低 #

np.add.at需要根据y的索引，逐个更新new_centers的对应行。由于y是随机生成的簇标签，内存访问是非连续的随机访问——每次更新的可能是完全不相邻的内存位置，这会导致CPU缓存频繁失效，大量时间浪费在等待内存数据加载上。

而简易实现中，X[y == i].sum(0)是连续读取X中属于同一簇的样本，内存访问模式是连续的，缓存命中率高，numpy对连续数组的求和操作又有专门的向量化优化，效率自然更高。

2. 通用实现的额外开销 #

np.add.at是一个通用的操作，需要处理各种复杂索引场景（比如重复索引、多维索引、广播等），这带来了大量的分支判断和参数校验开销。相比之下，简易实现的sum(0)是高度特化的向量化操作，底层没有多余的逻辑。

3. PyTorch的针对性优化 #

PyTorch的index_add_和index_put_是针对深度学习中频繁出现的分组累加场景设计的，尤其是你的版本带CUDA支持，直接利用了GPU的并行计算能力（即使CPU版本，PyTorch的底层实现也比numpy更偏向这类操作的优化）。

更高效的numpy替代方案 #

如果不想用PyTorch，可以试试用np.bincount实现，性能通常比前两种numpy方法更好：

# 按特征列分别用bincount求和，再转置得到结果
new_centers = np.array([np.bincount(y, weights=X[:, j], minlength=n_clusters) for j in range(n_features)]).T

这个方法利用bincount的向量化特性，避免了随机内存访问，同时没有Python循环的开销，测试下来性能会接近甚至超过简易实现。

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