【Pytorch】张量的维度/轴/dim的理解

原创

已于 2022-09-29 18:39:29 修改 · 3.2k 阅读

15 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #深度学习 #python

于 2022-09-29 16:04:46 首次发布

本文深入探讨了Pytorch中张量操作的关键概念——轴(axis/dim)，通过实例解析torch.cat()、torch.stack()、torch.sum()、torch.mean()和torch.softmax()等函数的轴参数使用。阐述了轴在不同张量运算中的作用，帮助读者掌握如何根据轴来处理多维数据。

1 前言

在深度学习中一旦涉及到三维或者更高维的数组（张量），Pytorch中各种涉及张量的函数操作都会通过参数axis/dim来设定操作的维度。很多张量的运算、神经网络的构建，都会涉及到轴，但到底取哪个轴，却不是那么容易把握。
下面会针对轴/axis/dim，基于 Pytorch 的代码和实例，尝试去理清张量运算中轴/axis/dim的设定。
以下内容为自学的记录，其中多有借鉴他人的博客，并在参考文献中一并给出引用。

2 轴的概念

对于一个张量，它的shape有几维，就对应有几个轴，也就对应着张量的层级。

import torch
a = torch.Tensor([[1,2,3], [4,5,6]]) 
# a shape [2, 3]
b = torch.Tensor([[7,8,9], [10,11,12]])
# b shape [2, 3]
c = torch.Tensor([[[1,2,3], [4,5