获取 2D PyTorch 张量对象的每一列的平均值

Question

我想知道什么是最简单的方法来获取矩阵每列的平均值，即我的张量只有两个维度，形状为 (m X n)。

例如，如果我有一个张量对象

T = torch.FloatTensor([[2.6, 5], [1.7, 6], [3.2, 7], [2.1, 8]])

我想要一些可以 return 张量对象的函数，如下所示

([2.4, 6.5])

我知道我可以使用 for 循环，但我想知道是否有 PyTorch 内置方法可以为我执行此操作，或者是否有更优雅的方法来实现此目的。

for i in range(T.size(1)):
  mean[i] = torch.mean(T[:, i])

感谢您的帮助。

Answer 1

请阅读 torch.mean 上的文档。它有一个可选参数 dim:

dim (int or tuple of ints) – the dimension or dimensions to reduce.

你的情况：

T.mean(dim=0)

会给你第一个暗淡的平均值。

请注意，pytorch（以及 numpy）中的许多其他“缩减”操作也有此选项以仅沿特定维度进行缩减（例如，std、sum 等.).

Get the mean of each column for a 2D PyTorch tensor object