【Pytorch进阶】torch.sum()中dim的使用

1 前言

  当开始用PyTorch张量做一些基本的运算时,如求和，它看起来很简单。但是对于多维的张量求和时就需要仔细思考一下(【Pytorch】张量的维度/轴/dim的理解)。

2 torch.sum()

2.1 一维求和

import torch

x = torch.tensor([1, 2, 3])
torch.sum(x)
tensor(6)

2.2 二维矩阵求和

import torch

x = torch.tensor(      
      [[1, 2, 3],      
      [4, 5, 6]]) 
x.shape  # torch.Size([2, 3])
print(torch.sum(x, dim=0))
# [5, 7, 9]
print(torch.sum(x, dim=1))
# [6, 15]

2.3 三维矩阵求和

y = torch.tensor([
     [[1, 2, 3],
      [4, 5, 6]],
     [[1, 2, 3],
      [4, 5, 6]],
     [[1, 2, 3],
      [4, 5, 6]]])
   
y.shape  # torch.Size([3, 2, 3])

torch.sum(y, dim=0) 
#tensor([[ 3,  6,  9],
#        [12, 15, 18]])

torch.sum(y, dim=1)  # torch.Size([3, 3])
#tensor([[5, 7, 9],
#        [5, 7, 9],
#        [5, 7, 9]])
y = torch.tensor([
    [[1, 2, 3],
     [4, 5, 6]],
    [[1, 2, 3],
     [4, 5, 7]],
    [[1, 2, 3],
     [4, 5, 8]]])
print(torch.sum(y, dim=2))
print(torch.sum(y, dim=2).shape)

#tensor([[ 6, 15],
#        [ 6, 16],
#        [ 6, 17]])
#torch.Size([3, 2])

图解可以看参考文献【1】

3 参考文献

[1]理解PyTorch中的dimensions维度，三维张量求和过程https://blog.csdn.net/crj0926/article/details/120587826