通过特定维度的三维张量的平均值计算的二维张量

Question

我有一个三维张量，形状为(3000，20，5(。我想创建一个二维张量，形状(3000，5(，使用三维张量的第二个维度的平均值。

所以基本上，我想做一些类似的事情：

mean_value = torch.mean(3d_tensor[0][:][0])

但是获取维度一到三的所有值的值。我可以做一个for循环，例如：

for j in range(0, 3d_tensor.size()[2]):
for i in range(0, len(3d_tensor)):
mean_values[j][i] = torch.mean(3d_tensor[i][:][j])

但这需要很长时间来处理大量数据。

Answer 1

您可以简单地指定应该取平均值的轴：

mean = torch.mean(tensor, dim=1)

这会给你形状(3000，5(的数据