PyTorch：torch.sum

torch.sum

函数定义：

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

作用：返回输入tensor的指定维度dim上的和。参数keepdim表示输出张量是否保留原有的dim。

应用举例：

例1——输出tensor单个维度的和

a = torch.randint(2, 10,(6,4))      # 创建shape为6*4，值为[2,10]的随机整数的tensor
b = torch.sum(a, dim=0)  # 返回a的第0维度(列)上的和
print('a:', a)
print('b:', b)

'''   运行结果   '''
a: tensor([[4, 2, 9, 6],
           [3, 2, 5, 2],
           [5, 6, 5, 4],
           [4, 4, 4, 7],
           [2, 4, 7, 5],
           [8, 4, 2, 8]])
b: tensor([26, 22, 32, 32])

例2——在多个维度上计算和

import torch
a = torch.randint(2, 10,(4, 5, 6))      # 创建shape为6*4，值为[2,10]的随机整数的tensor
b = torch.sum(a, (1, 2))  # 计算a的第1、2维度上的和，即结果分别是4个5*6的块的和
print('a:', a)
print('b:', b)

'''   运行结果   '''
a: tensor([[[5, 9, 7, 5, 4, 4],
            [6, 9, 6, 5, 5, 2],
            [4, 6, 9, 4, 8, 4],
            [5, 3, 3, 9, 3, 4],
            [7, 6, 8, 4, 5, 4]],    注：结果的第一个值163就是对该5*6的块进行求和，下同

           [[2, 9, 2, 6, 8, 3],
            [9, 7, 7, 5, 9, 5],
            [4, 7, 2, 2, 6, 2],
            [5, 9, 4, 8, 8, 6],
            [3, 2, 5, 3, 5, 6]],

           [[7, 2, 6, 3, 9, 8],
            [3, 3, 2, 3, 3, 7],
            [2, 9, 8, 2, 4, 9],
            [4, 8, 9, 3, 2, 6],
            [7, 2, 5, 2, 4, 9]],

           [[9, 6, 8, 5, 2, 7],
            [9, 8, 4, 7, 2, 7],
            [2, 9, 6, 7, 7, 9],
            [8, 5, 9, 7, 2, 2],
            [2, 5, 3, 3, 4, 2]]])
b: tensor([163, 159, 151, 166])

例3——保留原有的dim信息

a = torch.randint(2, 10,(6,4))      # 创建shape为6*4，值为[2,10]的随机整数的tensor
b = torch.sum(a, dim=0, keepdim=True)  # 返回a的第0维度(列)上的和，返回结果，并保留原始的维度信息
print('a:', a)
print('shape of a:', a.shape)
print('b:', b)  # 结果还是2维的
print('shape of b:', b.shape)

'''   运行结果   '''
a: tensor([[4, 4, 6, 7],
           [9, 8, 2, 7],
           [5, 9, 2, 2],
           [8, 8, 4, 5],
           [2, 3, 5, 8],
           [7, 6, 3, 3]])
shape of a: torch.Size([6, 4])
b: tensor([[35, 38, 22, 32]])
shape of b: torch.Size([1, 4])  注：keepdim默认为False，即默认结果中没有dim维度，这里保留了dim维度