举一个非常简单直观的例子:
首先,利用torch创建一个矩阵A:
A=torch.tensor([[1,2,3],[2,3,4]])
矩阵A的大小为:2*3
(1)dim=0:
B=torch.mean(A.float(),dim=0)
此时打印B为:
tensor([1.5000, 2.5000, 3.5000]) #说明dim=0时,是竖着计算均值,即按列计算;
(2) dim=1:
C=torch.mean(A.float(),dim=1)
此时打印C为:
tensor([2., 3.]) #说明dim=1时,是横向计算均值,即按行计算.