无需列表推导式，如何高效实现PyTorch张量指定区间批量求和？

核心思路是利用**前缀和（Prefix Sum）**将区间求和转化为张量级别的差值运算，完全替代Python循环/列表推导式，充分发挥PyTorch的并行计算能力。以下是具体实现步骤：

1. 计算counts的前缀和 #

counts是形状(num_features, num_bins)的张量，我们对每个特征的bins维度计算前缀和，得到(num_features, num_bins+1)的前缀和张量。这样任意闭区间[min_idx, max_idx]的元素和，等价于prefix_counts[k][max_idx+1] - prefix_counts[k][min_idx]，和你原来的counts[k][min:max+1].sum()结果完全一致。

# counts shape: (num_features, num_bins)
prefix_counts = torch.cat(
    [torch.zeros((counts.shape[0], 1), device=counts.device), counts.cumsum(dim=1)],
    dim=1
)
# prefix_counts shape: (num_features, num_bins + 1)

2. 批量计算所有区间的和 #

针对你的min/max张量（形状(num_samples, num_clusters, num_features)），通过张量维度置换和高级索引，一次性提取所有样本-簇-特征组合对应的前缀和差值：

# min/max shape: (num_samples, num_clusters, num_features)
max_plus_1 = max + 1

# 置换维度，让num_features和prefix_counts的第一维度对齐
sum_max = prefix_counts[torch.arange(counts.shape[0]), max_plus_1.permute(2,0,1)].permute(1,2,0)
sum_min = prefix_counts[torch.arange(counts.shape[0]), min.permute(2,0,1)].permute(1,2,0)

# 计算所有区间的和，结果形状: (num_samples, num_clusters, num_features)
interval_sum = sum_max - sum_min

3. 执行后续运算并求和 #

完成除以divisor、平方操作后，沿特征维度（dim=2）求和得到最终结果：

data_mass = ((interval_sum / divisor) ** 2).sum(dim=2)

完整示例验证（用你提供的测试数据） #

import torch

# 测试输入
min_tensor = torch.tensor([[[1, 2, 3, 1], 
                           [2, 1, 1, 2],
                           [1, 2, 2, 1]], 
                          [[2, 3, 2, 1], 
                           [3, 3, 1, 2],
                           [1, 0, 2, 1]]])
max_tensor = torch.tensor([[[3, 3, 4, 4], 
                           [3, 2, 3, 4],
                           [2, 4, 3, 2]], 
                          [[4, 4, 3, 3], 
                           [4, 4, 2, 3],
                           [2, 1, 3, 2]]])
counts = torch.tensor([[1, 2, 3, 4, 5], 
                       [2, 5, 3, 1, 1],
                       [1, 2, 3, 4, 5], 
                       [2, 5, 3, 1, 1]])
divisor = 10.0

# 步骤1：前缀和
prefix_counts = torch.cat([torch.zeros((counts.shape[0], 1)), counts.cumsum(dim=1)], dim=1)

# 步骤2：批量区间求和
max_plus_1 = max_tensor + 1
sum_max = prefix_counts[torch.arange(counts.shape[0]), max_plus_1.permute(2,0,1)].permute(1,2,0)
sum_min = prefix_counts[torch.arange(counts.shape[0]), min_tensor.permute(2,0,1)].permute(1,2,0)
interval_sum = sum_max - sum_min

# 步骤3：最终运算
data_mass = ((interval_sum / divisor) ** 2).sum(dim=2)

print(data_mass)

性能优势 #

• 完全消除Python层面的循环，所有操作都是PyTorch原生张量运算，能充分利用GPU并行计算（若使用CUDA）或CPU向量优化
• 针对你给出的实际尺寸（counts(16,100)、min/max(2708,7,16)），计算耗时可降至数秒级别，远优于原来的18-20秒

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