与其讨论到底是根据行还是列来划分,不如直接上例子。
二维情况下
创建一个张量
c=torch.arange(6).view(3,2)
dim=0时
tensor([[0, 1],
[2, 3],
[4, 5]])
tensor([6, 9])
dim=1时
tensor([[0, 1],
[2, 3],
[4, 5]])
tensor([1, 5, 9])
三维情况下
先创建一个三维张量,dim=0时
c=torch.arange(24).view(3,2,4)
torch.sum(c,dim=0)
tensor([[[ 0, 1, 2, 3],
[ 4, 5, 6, 7]],
[[ 8, 9, 10, 11],
[12, 13, 14, 15]],
[[16, 17, 18, 19],
[20, 21, 22, 23]]])
tensor([[24, 27, 30, 33],
[36, 39, 42, 45]])
当dim=1时
tensor([[[ 0, 1, 2, 3],
[ 4, 5, 6, 7]],
[[ 8, 9, 10, 11],
[12, 13, 14, 15]],
[[16, 17, 18, 19],
[20, 21, 22, 23]]])
tensor([[ 4, 6, 8, 10],
[20, 22, 24, 26],
[36, 38, 40, 42]])
dim=2时
tensor([[[ 0, 1, 2, 3],
[ 4, 5, 6, 7]],
[[ 8, 9, 10, 11],
[12, 13, 14, 15]],
[[16, 17, 18, 19],
[20, 21, 22, 23]]])
tensor([[ 6, 22],
[38, 54],
[70, 86]])