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

原创已于 2022-09-29 17:06:54 修改 · 2k 阅读

6 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #python #深度学习

于 2022-09-29 17:03:56 首次发布

Pytorch学习专栏收录该内容

32 篇文章

订阅专栏

本文介绍了PyTorch中torch.sum()函数的使用，包括一维、二维和三维矩阵的求和操作。通过示例展示了如何按不同维度进行求和，并提供了相关代码实现。对于理解和应用PyTorch张量运算具有指导意义。

部署运行你感兴趣的模型镜像

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])