torch.nn.functional.softmax(input) 维度学习

1. softmax（input， dim=none）

import torch.nn.functional as F
import torch
a = torch.randn((2,3,2))
b = F.softmax(a,1)
c = F.softmax(a,2)
d = F.softmax(a,0)
"""
a
Out[5]: 
tensor([[[ 1.8777,  0.8479],
         [-0.7809,  2.1866],
         [-0.5761,  0.6063]],

        [[ 1.5799, -0.7697],
         [-1.2000, -0.6929],
         [-0.2571, -0.9807]]])

b   #F.softmax(a,1)
Out[6]: 
tensor([[[0.8650, 0.1786],
         [0.0606, 0.6811],
         [0.0744, 0.1403]],

        [[0.8188, 0.3461],
         [0.0508, 0.3737],
         [0.1304, 0.2802]]])
c   #F.softmax(a,2)
Out[8]: 
tensor([[[0.7369, 0.2631],
         [0.0489, 0.9511],
         [0.2346, 0.7654]],

        [[0.9129, 0.0871],
         [0.3759, 0.6241],
         [0.6734, 0.3266]]])
d  #F.softmax(a,0)
Out[12]: 
tensor([[[0.5739, 0.8345],
         [0.6033, 0.9468],
         [0.4209, 0.8302]],

        [[0.4261, 0.1655],
         [0.3967, 0.0532],
         [0.5791, 0.1698]]])
"""

2 dim

a.size()
#Out[13]: torch.Size([2, 3, 2])

（1）dim=1时

a是3维结构，dim=0,1,2 代表三个维度的索引。dim=1时，在如图的方向上进行softmax运算。

（2）dim=2时

（3）dim=0时