如何在PyTorch/Numpy中并行按索引张量计算子张量局部最大值？

我们有形状为(1, n)的张量x，以及形状为(1, k)的索引张量d，需要计算以下k个分段子张量的最大值：

x[0:d[0]], x[d[0]:d[1]], x[d[1]:d[2]], ..., x[d[-2]: d[-1]]
最终输出形状为(1, k)的张量，包含各分段的最大值。for循环实现速度过慢，以下是PyTorch和Numpy的并行实现方案：

PyTorch 实现 #

推荐使用torch.scatter_reduce结合分段索引生成的方式，全程无显式循环，效率拉满：

import torch

# 示例输入
x = torch.randn(1, 10)  # shape (1, 10)
d = torch.tensor([[2, 5, 7, 10]])  # shape (1, 4)，注意最后一个索引需等于x的长度n

# 计算各分段长度
seg_lengths = torch.diff(torch.cat([torch.tensor([0], device=x.device), d[0]]))
# 生成每个元素对应的分段ID
seg_ids = torch.arange(d.shape[1], device=x.device).repeat_interleave(seg_lengths).unsqueeze(0)
# 并行计算各分段最大值
max_vals = torch.zeros(1, d.shape[1], device=x.device)
max_vals = max_vals.scatter_reduce(1, seg_ids, x, reduce='max', include_self=False)

print(max_vals)

Numpy 实现 #

最优方案是用np.maximum.reduceat，这是Numpy原生的分段归约函数，完全并行：

import numpy as np

# 示例输入
x = np.random.randn(1, 10)  # shape (1, 10)
d = np.array([[2, 5, 7, 10]])  # shape (1, 4)

# 构造分段起始索引：0, d[0], d[1], ..., d[k-2]
indices = np.concatenate([[0], d[0][:-1]])
# 并行计算各分段最大值并reshape为目标形状
max_vals = np.maximum.reduceat(x[0], indices).reshape(1, -1)

print(max_vals)

如果需要更直观的拆分方式，也可以用np.split（效率稍低但易读）：

segments = np.split(x[0], d[0][:-1])
max_vals = np.array([np.max(seg) for seg in segments]).reshape(1, -1)

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