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维度