PyTorch中如何高效实现tf.math.unsorted_segment_max功能？

针对你需要处理20万条数据且频繁执行的场景，循环实现效率过低，以下是两种高效的PyTorch解决方案：

方法1：使用torch.scatter_reduce（推荐，PyTorch 1.12+） #

PyTorch 1.12及以上版本原生提供了torch.scatter_reduce函数，完美对应TensorFlow的tf.math.unsorted_segment_max功能，底层经过并行优化，性能最优。

代码示例：

import torch

data = torch.tensor([5, 1, 7, 2, 3, 4], dtype=torch.float32)
segment_idxs = torch.tensor([0, 0, 1, 1, 0, 1], dtype=torch.int64)
num_segments = 2

# 初始化结果为负无穷，确保所有数据值都能覆盖初始值
max_at_idxs = torch.full((num_segments,), float('-inf'), dtype=data.dtype)
# 执行分组max聚合
max_at_idxs = max_at_idxs.scatter_reduce(0, segment_idxs, data, reduce='max', include_self=False)

print('max_at_idxs: ')
print(max_at_idxs)

输出：

max_at_idxs: 
tensor([5., 7.])

参数说明：

• 0：指定在第0维度执行scatter操作
• segment_idxs：分组索引数组
• data：待聚合的原始数据
• reduce='max'：设置聚合方式为取最大值
• include_self=False：确保初始的负无穷值不参与和自身的比较，只使用data中的有效数据

方法2：兼容旧版本PyTorch的实现（PyTorch <1.12） #

如果无法升级PyTorch版本，可以通过排序+分段取max的方式实现，效率远高于逐元素循环：

import torch

data = torch.tensor([5, 1, 7, 2, 3, 4], dtype=torch.float32)
segment_idxs = torch.tensor([0, 0, 1, 1, 0, 1], dtype=torch.int64)
num_segments = 2

# 按分组索引排序，将同一组的数据集中
sorted_indices = torch.argsort(segment_idxs)
sorted_segments = segment_idxs[sorted_indices]
sorted_data = data[sorted_indices]

# 计算每个分组的边界位置
_, group_counts = torch.unique(sorted_segments, return_counts=True)
group_boundaries = torch.cumsum(group_counts, dim=0)
group_boundaries = torch.cat([torch.tensor([0]), group_boundaries])

# 分段提取每组的最大值
max_at_idxs = torch.zeros(num_segments, dtype=data.dtype)
for i in range(num_segments):
    start = group_boundaries[i]
    end = group_boundaries[i+1]
    max_at_idxs[i] = sorted_data[start:end].max()

print('max_at_idxs: ')
print(max_at_idxs)

输出：

max_at_idxs: 
tensor([5., 7.])

该方法通过排序将同组数据集中，再利用PyTorch的向量化max操作替代循环，大幅提升执行效率。

性能对比 #

针对20万条数据的场景：

• 循环实现：单轮执行时间通常在几十毫秒到上百毫秒（CPU/GPU差异较大）
• scatter_reduce实现：单轮执行时间在微秒级（GPU上几乎可忽略）
• 排序分段实现：单轮执行时间在几毫秒（GPU环境下）

优先推荐使用torch.scatter_reduce，建议升级PyTorch版本以获得最优性能。

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