eniops库中reduce函数使用方法

安安爸Chris2025-01-26 22:03

reduce 是 eniops 中的一个常用函数，用于对张量进行降维操作。它允许你通过指定维度名称和操作类型（如求和、均值等）来简化张量的形状。

python 复制代码

import eniops
import torch

# 创建一个示例张量
x = torch.randn(2, 3, 4)

# 使用 reduce 进行降维操作
result = eniops.reduce(x, 'b c h -> b h', 'mean')

print(result.shape)  # 输出: torch.Size([2, 4])

输入张量 x 的形状为 (2, 3, 4)，对应模式 'b c h'。

reduce 操作将 c 维度通过 'mean' 操作降维，最终输出形状为 (2, 4)，对应模式 'b h'。

除了mean，还有sum,max等降维方式.

如下，

python 复制代码

result = eniops.reduce(x, 'b c h -> b h', 'sum')
print(result.shape)  # 输出: torch.Size([2, 4])

python 复制代码

result = eniops.reduce(x, 'b c h -> b h', 'max')
print(result.shape)  # 输出: torch.Size([2, 4])