©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，但它们含义是一样的。

轴的使用

在做张量的拼接操作时，axis/dim 设定了哪个轴，那对应的轴在拼接之后张量数会发生变化

>> torch.cat((a,b), dim=0)
tensor([[ 1.,  2.,  3.],[ 4.,  5.,  6.],[ 7.,  8.,  9.],[10., 11., 12.]])>> torch.cat((a,b), dim=1)
tensor([[ 1.,  2.,  3.,  7.,  8.,  9.],[ 4.,  5.,  6., 10., 11., 12.]])

对于上面 torch 中的 cat 操作，当设定 dim=0 时，两个维度是 (2,3) 的张量合并成了一个 (4,3) 的张量，在第 0 维，张量数从 2 变成了 4，第 1 维没有变化；当设定 dim=1 时，在第 1 维，张量数从 3 变成了 6，第 0 维没有变化。

在做张量的运算操作时，axis/dim 设定了哪个轴，就会遍历这个轴去做运算，其他轴顺序不变。

>> torch.softmax(a, dim=0)
tensor([[0.0474, 0.0474, 0.0474],[0.9526, 0.9526, 0.9526]])>> torch.softmax(a, dim=1)
tensor([[0.0900, 0.2447, 0.6652],[0.0900, 0.2447, 0.6652]])

对于上面 torch 中的 softmax 操作，当设定 dim=0 时，就是其他轴不变，单次遍历 dim=0 轴的所有元素去做运算，上例中就相当于分别取了张量 a 中的第 0 列、第 1 列、第 2 列去做计算。

换一个角度，假设用 for 循环去遍历一个张量，那运算中设定的 dim 就是被放在最内层的 for 循环，其它的轴保持正常的顺序。

可以用下面的例子作为验证，这里 tensor  c  的 shape 是 (m,n,p)，用 for 循环去计算 torch.softmax(c, dim=1) 。

# for循环计算方式
c = torch.Tensor([[[1,2,3], [4,5,6]], [[7,8,9], [10,11,12]]])    # shape (2,2,3)
m,n,p = c.shape
res = torch.zeros((m,n,p))
for i in range(m):for j in range(p):res[i,:,j] = torch.softmax(torch.tensor([c[i,k,j] for k in range(n)]), dim=0)  #这里对应最内层的for循环# 库函数设定轴计算方式
res1 = torch.softmax(c, dim=1)
print(res.equal(res1))        # True

axis/dim 使用小总结：

1. 在做张量的拼接操作时，axis/dim 设定了哪个轴，那对应的轴在拼接之后樟树会发生变化 axis/dim 设定了哪个轴，那对应的轴在拼接之后的张量数会发生变化

2. 在做张量的运算操作时，axis/dim 设定了哪个轴，就会遍历这个轴去做运算，其他轴顺序不变

实际上，第一条拼接操作也可以用第二条去理解，但拼接的轴张量数会发生变化更好理解和记忆。

轴的实例

其实一个轴设定的变化，会带来很大的差异，最典型的就是 BatchNorm 和 LayerNorm 了。

BatchNorm 和 LayerNorm 是针对数据的不同轴去做 norm，假设输入数据的维度是（N,C,H,W），分别对应 batch 数，核数，高，宽，BatchNorm 就对应 dim=0，LayerNorm 就对应 dim=1，在不考虑移动平均等具体细节问题时，两者在形式上可以统一，只有一个 dim 参数的差别。

Pytorch 的实现（简化版）如下：

class Norm(nn.Module):def __init__(self, num_features, variance_epsilon=1e-12):super(Norm, self).__init__()self.gamma = nn.Parameter(torch.ones(num_features))self.beta = nn.Parameter(torch.zeros(num_features))self.variance_epsilon = variance_epsilon    # 一个很小的常数，防止除0def forward(self, x, dim):u = x.mean(dim, keepdim=True)s = (x - u).pow(2).mean(dim, keepdim=True)x_norm = (x - u) / torch.sqrt(s + self.variance_epsilon)return self.gamma * x_norm + self.beta

当然，不仅仅是在深度学习里面，在 Numpy，Pandas中，轴的设定都经常会遇到，但它们都是相通的，希望本文能帮你更好的理解它 —> 轴。

更多阅读

#投 稿 通 道#

 让你的论文被更多人看到 

如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。

总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。

PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学习心得或技术干货。我们的目的只有一个，让知识真正流动起来。

???? 来稿标准：

• 稿件确系个人原创作品，来稿需注明作者个人信息（姓名+学校/工作单位+学历/职位+研究方向）

• 如果文章并非首发，请在投稿时提醒并附上所有已发布链接

• PaperWeekly 默认每篇文章都是首发，均会添加“原创”标志

???? 投稿邮箱：

• 投稿邮箱：hr@paperweekly.site

• 所有文章配图，请单独在附件中发送

• 请留下即时联系方式（微信或手机），以便我们在编辑发布时和作者沟通

????

现在，在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧

关于PaperWeekly

PaperWeekly 是一个推荐、解读、讨论、报道人工智能前沿论文成果的学术平台。如果你研究或从事 AI 领域，欢迎在公众号后台点击「交流群」，小助手将把你带入 PaperWeekly 的交流群里。