关于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则是去掉那一维的个数。

相关推荐
风舞雪凌月3 分钟前
【总结】国产AI大模型公司汇总
人工智能
Hali_Botebie4 分钟前
【光流】自动驾驶光流任务 DeFlow: Decoder of Scene Flow Network in Autonomous Driving
人工智能·机器学习·自动驾驶
IT_陈寒7 分钟前
被Vite的HMR坑惨了,原来这样配置才能用对!
前端·人工智能·后端
“码”力全开10 分钟前
解耦安防碎片化:基于 Docker 与边缘计算的 AI 视频中台架构设计(支持 GB28181/RTSP 与源码交付)
人工智能·docker·边缘计算
sali-tec12 分钟前
C# 基于OpenCv的视觉工作流-章80-长短脚
图像处理·人工智能·opencv·算法·计算机视觉
AI科技星12 分钟前
国家重点研发计划项目申报书
人工智能·线性代数·架构·概率论·学习方法
智联视频超融合平台13 分钟前
AI赋能传统电厂:2025能源革命的智慧引擎
人工智能·能源
qcx2320 分钟前
【系统学AI】23 AI 时代产品运营与获客全景:CRM SaaS 大变局 + 增长新范式(2026 调研报告)
人工智能·产品运营·产品设计·ai agent·ai native
白狐_79821 分钟前
从功能开发到开源维护:一个 Python 可视化项目的 Git 分支、维护文件与 PR 流程实践
git·python·开源
叶修_A21 分钟前
【COZE-08】Prompt工程进阶 - 结构化输出与思维链
大数据·人工智能·prompt