【PyTorch】torch.mean(), dim=0, dim=1 详解

最新推荐文章于 2024-03-21 10:50:23 发布

shuaiqidexiaojiejie

最新推荐文章于 2024-03-21 10:50:23 发布

阅读量502

点赞数

分类专栏： # torch 文章标签：深度学习

本文链接：https://blog.youkuaiyun.com/shuaiqidexiaojiejie/article/details/124384985

版权

本文详细解释了在PyTorch中如何使用torch.mean()函数进行维度操作，包括dim=0和dim=1时的效果。通过实例展示了不同维度选择下，计算平均值如何影响数据的形状，并提到了数据类型的要求。同时提到，可以使用keepdims=True保持维度不变。这些概念同样适用于其他类似函数如torch.sum()。

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

创建一个tensor,这个tensor是一个元素类型为浮点型的2维数组

import torch

s = torch.arange(6,dtype=float).reshape((2,3))
print(s)
print(s.shape)# 查看tensor的形状

tensor([[0., 1., 2.],
        [3., 4., 5.]], dtype=torch.float64)
torch.Size([2, 3])

dim属性的全称是dimension，表示维度。dim=0为第0个维度，代表行。

对于torch.mean(s,dim=0),表示跨行求平均。

得到的结果是一个向量，分别对应于 1.5=(0.0+3.0)/2, 2.5=(1.0+4.0)/2, 3.5=(2.0+5.0)/2

s1 = torch.mean(s, dim=0)
print(s1)

tensor([1.5000, 2.5000, 3.5000], dtype=torch.float64)

同理，对于dim=1为第一个维度，代表列。

对于torch.sum(s,dim=1),表示跨列求平均。

得到的结果同样是一个向量，分别对应于 1.0=(0.0+1.0+2.0)/3， 4.0=(3.0+4.0+5.0)/3

s2 = torch.mean(s, dim=1)
print(s2)

tensor([1., 4.], dtype=torch

最低0.47元/天解锁文章