pytorch中选取不同维度sum和mean方法理解

最新推荐文章于 2025-03-24 11:07:08 发布

木禾DING

最新推荐文章于 2025-03-24 11:07:08 发布

阅读量1.4k

点赞数 3

CC 4.0 BY-SA版权

文章标签： pytorch 人工智能 python

本文链接：https://blog.youkuaiyun.com/ding_programmer/article/details/131572257

PyTorch的sum()函数用于对张量进行求和操作。它可以对所有元素求和得到标量，或者按指定维度求和得到新的张量。例如，对于一个(2,3,4)的矩阵，sum(dim=0)会得到(3,4)的张量，sum(dim=1)得到(2,4)的张量，sum(dim=2)或sum(dim=-1)得到(2,3)的张量。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

在PyTorch中，sum()函数用于对输入张量的所有元素进行求和操作。该函数的语法如下：

torch.sum(input, dim=None, keepdim=False, dtype=None)

具体而言，sum()函数会对输入张量的所有元素进行求和操作，并返回一个标量值。

如果指定了dim参数，则会沿着指定的维度对输入张量进行求和操作，并返回一个形状与输入张量除了指定维度之外的维度相同的张量。

如果指定了keepdim参数，则返回的张量会保持与输入张量相同的维度数，并且在指定维度上的大小为1。

具体算维度，采用消解法，也就是 dim=0，那就是消去第0维，dim=1，消去第1维

比如一个矩阵维度(2,3,4)，dim=0，那就是消去第0维，变成了 (3,4)，消去第1维，变成了 (2,4)，消去第2维，变成了 (2,3)，dim=-1，也就是最后一维，在这个矩阵中也就是第二维。

import torch

x = torch.arange(24).view(2,3,4)
print(x)

tensor([[[ 0,  1,  2,  3],
         [ 4,  5,  6,  7],
         [ 8,  9, 10, 11]],

        [[12, 13, 14, 15],
         [16, 17, 18, 19],
         [20, 21, 22, 23]]])


y = torch.sum(x)
print(y)

tensor(276)


y = torch.sum(x,dim=0)
print(y)  
print(y.shape)


tensor([[12, 14, 16, 18],
        [20, 22, 24, 26],
        [28, 30, 32, 34]])
torch.Size([3, 4])





y = torch.sum(x, dim=1)
print(y)  
print(y.shape)

tensor([[12, 15, 18, 21],
        [48, 51, 54, 57]])
torch.Size([2, 4])


y = torch.sum(x, dim=2)
print(y)  
print(y.shape)

tensor([[ 6, 22, 38],
        [54, 70, 86]])
torch.Size([2, 3])

y = torch.sum(x, dim=-1)
print(y)  
print(y.shape)

tensor([[ 6, 22, 38],
        [54, 70, 86]])
torch.Size([2, 3])