首先定义张量A
A = torch.arange(20, dtype=torch.float32).reshape(5, 4)
tensor([[ 0., 1., 2., 3.],[ 4., 5., 6., 7.],[ 8., 9., 10., 11.],[12., 13., 14., 15.],[16., 17., 18., 19.]])
1. 降维求和
降维求和会沿指定轴降低张量的维度,使它变为一个标量。
A_sum_axis0 = A.sum(axis=0) # 压缩为一行
tensor([40., 45., 50., 55.]A_sum_axis1 = A.sum(axis=1) # 压缩为一列
tensor([ 6., 22., 38., 54., 70.]A_sum = A.sum(axis=[0, 1]) # 结果与 A.sum() 相同
tensor(190.)
2. 非降维求和
保持轴数不变
A_sum_axis0 = A.sum(axis=0, keepdims=True)
tensor([[40., 45., 50., 55.]])A_sum_axis1 = A.sum(axis=1, keepdims=True)
tensor([[ 6.],[22.],[38.],[54.],[70.]])A_sum = A.sum(axis=[0, 1], keepdims=True)
tensor([[190.]])
3. 累积求和
沿某个轴计算A元素的累积总和,此函数不会沿任何轴降低输入张量的维度。
A_sum_axis0 = A.cumsum(axis=0)
tensor([[ 0., 1., 2., 3.],[ 4., 6., 8., 10.],[12., 15., 18., 21.],[24., 28., 32., 36.],[40., 45., 50., 55.]])A_sum_axis1 = A.cumsum(axis=1)
tensor([[ 0., 1., 3., 6.],[ 4., 9., 15., 22.],[ 8., 17., 27., 38.],[12., 25., 39., 54.],[16., 33., 51., 70.]])