如何将PyTorch张量按比例聚合为更小尺寸？

你可以通过几种简洁的方式实现这个需求，以下是常见的实用方法：

方法1：Reshape + 聚合函数（最直观高效） #

先将1D张量重塑为2D结构（每组占一行），再对每行应用求和、均值等聚合函数。这种方法要求原张量长度是分组大小的整数倍。

import torch

x = torch.tensor([1,2,3,4,5,6,7,8,9])
group_size = 3

# 重塑为(组数, 每组元素数)的2D张量
x_2d = x.reshape(-1, group_size)
# 按行求和
sum_tensor = x_2d.sum(dim=1)
print(sum_tensor)  # 输出: tensor([ 6, 15, 24])

# 按行求均值
mean_tensor = x_2d.mean(dim=1)
print(mean_tensor)  # 输出: tensor([2., 5., 8.])

方法2：使用torch.unfold #

unfold可以按指定窗口大小和步长提取子张量，将步长设为分组大小即可实现不重叠分组：

x_unfolded = x.unfold(dimension=0, size=group_size, step=group_size)
sum_tensor = x_unfolded.sum(dim=1)
print(sum_tensor)  # 输出: tensor([ 6, 15, 24])

方法3：torch.split + 批量聚合（灵活处理非整数倍长度） #

如果原张量长度不能被分组大小整除，split会自动将最后一组设为剩余元素，再通过列表推导式完成聚合：

# 示例：处理长度无法被分组大小整除的情况
x = torch.tensor([1,2,3,4,5,6,7])
groups = torch.split(x, group_size)
sum_tensor = torch.tensor([g.sum() for g in groups])
print(sum_tensor)  # 输出: tensor([ 6, 15,  7])

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