线性代数 动手学深度学习 pytorch

axis = 0按照某个维度求和，就是去掉了某个维度，比如 shape: [5, 4] axis = 0, 相当于去掉维度0，也就是5，剩下维度4 sum: [4].

deepdims=True 表示保持维度，但是去掉的维度axis 都为1.
比如 shape: [5, 4] axis = 0, deepdims=True , 结果就是sum: [1, 4]

例子：

import torch

a = torch.ones((2, 5, 4))
a.shape

torch.Size([2, 5, 4])

a.sum().shape

torch.Size([])

a.sum(axis=1)

tensor([[5., 5., 5., 5.],
        [5., 5., 5., 5.]])

a.sum(axis=1).shape

torch.Size([2, 4])

a.sum(axis=0).shape

torch.Size([5, 4])

a.sum(axis=0)

tensor([[2., 2., 2., 2.],
        [2., 2., 2., 2.],
        [2., 2., 2., 2.],
        [2., 2., 2., 2.],
        [2., 2., 2., 2.]])

a.sum(axis=[0,2]).shape

torch.Size([5])

a.sum(axis=[0,2])

tensor([8., 8., 8., 8., 8.])

a.sum(axis=1,keepdims=True).shape

torch.Size([2, 1, 4])

a.sum(axis=1,keepdims=True)

tensor([[[5., 5., 5., 5.]],

        [[5., 5., 5., 5.]]])

a.sum(axis=[0,2],keepdims=True).shape

torch.Size([1, 5, 1])

a.sum(axis=[0,2],keepdims=True)

tensor([[[8.],
         [8.],
         [8.],
         [8.],
         [8.]]])

参考

https://www.bilibili.com/video/BV1eK4y1U7Qy?share_source=copy_web