pytorch | softmax(x,dim=-1)参数dim的理解

转自：https://zhuanlan.zhihu.com/p/525276061

首先说一下Softmax函数，公式如下:

1. 三维tensor(C,H,W)

一般会设置成dim=0,1,2,-1的情况(可理解为维度索引)。其中2与-1等价，相同效果。

用一张图片来更好理解这个参数dim数值变化：

当dim=0时， 是对每一维度相同位置的数值进行softmax运算，和为1
当dim=1时， 是对某一维度的列进行softmax运算，和为1
当dim=2时， 是对某一维度的行进行softmax运算，和为1

准备工作：先随机生成一个（2,5,4）的矩阵，即两个维度的（5,4）矩阵

import torch 
import torch.nn.functional as F 
input= torch.randn(2,2,3))
print(input)

随机3维矩阵

(1) dim=0

dim=0

(2) dim=1

dim=1

(3) dim=2 或dim=-1

dim=2

dim=-1

2. 四维tensor(B,C,H,W)

是三维tensor的推广，其实三维tensor也可以是batchsize=1的四维tensor，只是dim的索引需要加1.

dim取值0,1,2,3,-1

准备工作：先随机生成一个(2,2,5,4)矩阵。其实随着dim增加（从0到3），相当于一层层剥开。

(1) dim=0

这时的视野应该放在整个tensor，每个batch(不同B)对应位置(相同CHW)求softmax

(2) dim=1

这时向里剥，每小块(不同C)对应位置(相同BHW)求softmax。

(3) dim=2

继续向里剥，每小块(不同H)对应位置(相同BCW)求Softmax。

(4) dim=3 或dim=-1

继续向里剥，也是最后一次。每个小块(不同W)对应元素(相同BCH)求softmax。