代码+实例：深度学习中的“轴”全解

©PaperWeekly 原创 · 作者｜海晨威

学校｜同济大学硕士生

研究方向｜自然语言处理

在深度学习中，轴，指的就是张量的层级，一般通过参数 axis/dim 来设定。很多张量的运算、神经网络的构建，都会涉及到轴，但到底取哪个轴，却不是那么容易把握。

下面会针对轴/axis/dim，基于 PyTorch 的代码和实例，尝试去理清张量运算中轴/axis/dim 的设定。

轴的概念

对于一个张量，它的 shape 有几维，就对应有几个轴，也就对应着张量的层级，最直观的可以通过看最前面的方括号数量来判断。

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

# torch.Size([2, 3])

上面的张量 a 和 b，都对应两个轴。axis/dim=0 对应 shape [2, 3] 中的 2，axis/dim=1 对应 shape [2, 3] 中的 3，而张量 c 有三个轴。

张量运算时对轴参数的设定非常常见，在 Numpy 中一般是参数 axis，在 Pytorch 中一般是参数 dim，但它们含义是一样的。