关于pytorch中的dim的理解

今天碰到一个代码看起来很简单,但是细究原理又感觉好像不太通不太对劲,就是多维tensor数据的操作,比如:y.sum(dim=2),乍一看很简单数据相加操作,但是仔细一想,这里在第3维度的数据到底是横向相加还是纵向相加,带着疑问实验几次就明白了。

首先给个完整的例子:

python 复制代码
import  torch

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

print(y.sum(dim=2))

这里的y.shape = (3, 2, 3),三个维度的数据,所以dim可以是0~2也可以是-1~-3。我们每个维度都进行操作一遍就清楚了。

  • 当dim=0时,相当于有3个二维的向量进行相加,结果还是一个二维向量(对应位置相加):

    y.shape = (3, 2, 3) ---> y.shape = (2, 3)
  • 当dim=1时,相当于有2个一维的向量进行相加×3,结果是1个一维向量×3则还是一个二维向量:

    y.shape = (3, 2, 3) ---> y.shape = (3, 3)
  • 当dim=2时,相当于有3个数值进行相加×2×3,结果两个值组成一维向量,三个一维向量组成二维向量:

    y.shape = (3, 2, 3) ---> y.shape = (3, 2)

其他的数据操作也是这样类似的思想。

总结:从中可以看出只要对一个n维度的数据的其中一维进行操作的话,得到的结果会是n-1维的向量,shape则是去掉那一维的个数。

相关推荐
阿坡RPA5 小时前
手搓MCP客户端&服务端:从零到实战极速了解MCP是什么?
人工智能·aigc
用户27784491049936 小时前
借助DeepSeek智能生成测试用例:从提示词到Excel表格的全流程实践
人工智能·python
机器之心6 小时前
刚刚,DeepSeek公布推理时Scaling新论文,R2要来了?
人工智能
算AI8 小时前
人工智能+牙科:临床应用中的几个问题
人工智能·算法
JavaEdge在掘金8 小时前
ssl.SSLCertVerificationError报错解决方案
python
我不会编程5558 小时前
Python Cookbook-5.1 对字典排序
开发语言·数据结构·python
凯子坚持 c9 小时前
基于飞桨框架3.0本地DeepSeek-R1蒸馏版部署实战
人工智能·paddlepaddle
老歌老听老掉牙9 小时前
平面旋转与交线投影夹角计算
python·线性代数·平面·sympy
满怀10159 小时前
Python入门(7):模块
python
无名之逆9 小时前
Rust 开发提效神器:lombok-macros 宏库
服务器·开发语言·前端·数据库·后端·python·rust