Pytorch tensor求和( tensor.sum())

时间:2024-01-26 21:00:45

1. torch.sum(input, dim, out=None)

参数说明:

input:输入的tensor矩阵。

dim:求和的方向。若input为2维tensor矩阵,dim=0,对列求和;dim=1,对行求和。注意:输入的形状可为其他维度(3维或4维),可根据dim设定对相应的维度求和。

out: 输出,一般不指定。

2. 例子

这里给了一个四维的tensor,代表的是输入到网络中的一张图片(batchSize,channel,height,width)。dim=1,对通道求和。

2.1 Code

 

 1 import torch
 2 
 3 
 4 def function():
 5     data1 = torch.rand([1, 3, 3, 3])
 6     print("data1_shape: ", data1.shape)
 7     print("data1: ", data1)
 8 
 9     data2 = torch.sum(data1, dim=1)
10     print("data2_shape: ", data2.shape)
11     print("data2: ", data2)
12 
13 
14 if __name__ == \'__main__\':
15     function()
View Code

 

 2.2 结果显示

对第二个通道求和的,因此第二个通道变为了1,默认被隐藏,不被显示。理论上的维度为[1, 1, 3, 3],隐藏后维度为[1, 3, 3]。