pytorch 中 torch.sum函数

 

sum 函数有两个定义：

torch.sum(input, *, dtype=None) → Tensor
torch.sum(input, dim, keepdim=False, *, dtype=None) → Tensor

** sum 函数本质上是对张量自身做自加操作，只是它可以定义，在不同维度上的相加，指定相加的维度会被降解掉。除非keepdim=True字段。**

 

In [2]:

import torch
from IPython.core.interactiveshell import InteractiveShell
InteractiveShell.ast_node_interactivity = "all"

a = torch.tensor([[1, 2, 3],[4,5,6]])
a.sum()  # 返回标量类型总和数据

 

Out[2]:

tensor(21)

 

In [14]:

a = torch.tensor([[[1, 2, 3],[4,5,6]],[[1, 2, 3],[4,5,6]]])
a.sum(dim=0)  # 对求和的部分，维度会下降
a.sum(dim=-1, keepdim=True)  # 将会保持维度不变

 

Out[14]:

tensor([[ 2,  4,  6],
        [ 8, 10, 12]])

Out[14]:

tensor([[[ 6],
         [15]],

        [[ 6],
         [15]]])