计算除第n维以外的所有维度的平均值

Question

我有一个形状为 [7, 12, 12, 197, 197] 的 PyTorch 张量，我想计算其平均值，使得结果向量的形状变为 [7,1,1,197,1]。除了使用 for 循环外，还有没有其他方法？能否给出具有执行时间基准的答案？

user370305 · Answer

这是使用 torch.mean 的简单方法，请检查是否有效。

import torch

tensor = torch.randn(4, 5, 6, 7)  # Example tensor
nth_dim = 2  # Dimension to keep

means = torch.mean(tensor, dim=torch.arange(tensor.ndim) != nth_dim)
print(means.shape)  # Output: torch.Size([4, 6, 7])

Nick Craver · Answer

在使用 numpy 库时，你可以通过给定一个包含用于聚合的维度/轴的元组，并设置 keepdims=True 参数来保持输出形状，从而调用 mean 函数：

out = arr.mean((1, 2, 4), keepdims=True)

示例：

np.random.seed(0)  # 设置随机种子确保可复现结果
arr = np.random.random((7, 12, 12, 197, 197))  # 创建一个五维数组
print(arr.shape)

out = arr.mean((1, 2, 4), keepdims=True)  # 在第 1、2、4 个维度上求平均并保持维度
print(out.shape)

输出：

(7, 12, 12, 197, 197)  # 输入数组的原始形状
(7, 1, 1, 197, 1)  # 求平均后保留维度的输出数组形状

Taryn · Answer

PyTorch 提供了一个带有 dim 参数的 mean 函数，该参数可以接受整数元组。

import torch
# 创建一个四维张量
tensor = torch.randn(4, 5, 6, 7)

# 定义要减少的维度
dims_to_reduce = (1, 3)

# 计算指定维度上的均值，并保持维度大小为1
means = torch.mean(tensor, dim=dims_to_reduce, keepdim=True)

# 打印计算后的均值张量的形状
print(means.shape)
# 输出: torch.Size([4, 1, 6, 1])

Vittorio Romeo · Answer

可以通过调用 mean 函数并指示其在哪些轴上进行平均计算，并且在计算过程中保持维度不变来完成这一操作。在 PyTorch 中实现这一功能的代码与在 NumPy 中非常相似，只不过将 axis 替换为 dim，并将 keepdims 替换为 keepdim。

以下是 PyTorch 中的例子：

import torch

# 创建一个五维张量
a = torch.rand(7, 12, 12, 197, 197)
print(a.shape)  # 输出：torch.Size([7, 12, 12, 197, 197])

# 指定轴并在保持维度的情况下计算平均值
b = torch.mean(a, dim=(1, 2, 4), keepdim=True)
print(b.shape)  # 输出：torch.Size([7, 1, 1, 197, 1])

而在 NumPy 中的相应实现如下：

import numpy as np

# 创建一个默认随机数生成器并生成五维数组
rng = np.random.default_rng(42)
a = rng.random((7, 12, 12, 197, 197))
print(a.shape)  # 输出：(7, 12, 12, 197, 197)

# 按指定轴计算平均值并在保持维度的前提下返回结果
b = np.mean(a, axis=(1, 2, 4), keepdims=True)
print(b.shape)  # 输出：(7, 1, 1, 197, 1)