如何实现有向图求和聚合操作的向量化优化？

核心问题说明 #

重复索引下的累加操作无法用普通切片赋值实现，需要使用支持无缓冲累加的专用API，以下是三种不同场景的实现思路：

思路1：纯NumPy实现（无额外依赖） #

利用np.add.at的原地无缓冲累加特性，完美处理重复索引的累加需求，代码示例如下：

import numpy as np

# 构造各维度对应索引
batch_idx = np.repeat(np.arange(batchSize), 2 * nEdges)
edge_seq = np.concatenate([np.arange(2 * e) for e in nEdges])
recv_flat = receivers[batch_idx, edge_seq]
feat_flat = localFeature[batch_idx, edge_seq, :2]

# 原地累加，自动处理重复索引
np.add.at(localSum, (batch_idx, recv_flat, slice(0, 2)), feat_flat)

该实现完全替换两层Python循环，底层由C实现，速度比原生循环高1~2个数量级。

思路2：深度学习框架实现（支持GPU加速） #

如果使用PyTorch/TensorFlow框架，可直接调用内置的scatter_add接口，还能利用GPU做并行加速，PyTorch示例代码如下：

import torch

batch_idx = torch.repeat_interleave(torch.arange(batchSize), 2 * nEdges)
edge_seq = torch.cat([torch.arange(2 * e) for e in nEdges])
recv_flat = receivers[batch_idx, edge_seq].long()
feat_flat = localFeature[batch_idx, edge_seq, :2]

localSum = torch.zeros(batchSize, max_node_num, 2, dtype=feat_flat.dtype)
localSum.scatter_add_(
    dim=1,
    index=recv_flat[:, None, None].expand(-1, 1, 2),
    src=feat_flat[:, None, :]
)

该实现可直接迁移到GPU运行，适合大图、大batch的高吞吐场景。

思路3：稀疏矩阵实现（兼容Matlab） #

如果需要适配Matlab环境，可使用稀疏矩阵乘法的通用逻辑，Python和Matlab都支持该思路：

• 对每个batch构造大小为(max_node_num, 2 * nEdges[j])的稀疏邻接矩阵，第receivers[j,i]行、第i列取值为1
• 用该稀疏矩阵左乘当前batch的边特征矩阵localFeature[j, :2*nEdges[j], :2]，直接得到对应batch的聚合结果localSum[j, :, :2]
  Matlab对稀疏矩阵运算的原生优化程度很高，该方案在Matlab上的运行效率远高于手写循环。

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