比如alpha = F.softmax(alpha, dim = 2),dim参数表示在alpha的第几个维度上进行softmax。
举例说明:
1. 随机生成一个维度为(2,2,3)的tensor
import torch.nn as nn
input = torch.randn(2, 2, 3)
print(input)
Out:
tensor([[[ 0.9235, -0.3383, -0.1117],
[ 0.5047, 2.6545, 0.5707]],
[[-1.5672, -0.7993, -0.8975],
[-0.7288, -0.4013, 1.1321]]])
m = nn.Softmax(dim=0)
n = nn.Softmax(dim=1)
k = nn.Softmax(dim=2)
print(m(input))
print(n(input))
print(k(input))
2. dim=0
Out:
tensor([[[0.9235, 0.6133, 0.6869],
[0.7744, 0.9550, 0.3632]],
[[0.0765, 0.3867, 0.3131],
[0.2256, 0.0450, 0.6368]]])
在第一个维度上进行softmax:
0.9235+0.0765=1,
0.6133+0.3867=1,
0.6869+0.3131=1,
0.7744+0.2256=1,
...
3. dim=1
Out:
tensor([[[0.6032, 0.0478, 0.3357],
[0.3968, 0.9522, 0.6643]],
[[0.3019, 0.4018, 0.1161],
[0.6981, 0.5982, 0.8839]]])
在第二个维度上进行softmax:
0.6032+0.3968=1,
0.3019+0.6981=1,
0.0478+0.9522=1,
...
4. dim=2
Out:
tensor([[[0.6104, 0.1728, 0.2168],
[0.0939, 0.8058, 0.1003]],
[[0.1957, 0.4219, 0.3824],
[0.1134, 0.1574, 0.7292]]])
在第三个维度上进行softmax:
0.6104+0,1728+0.2168=1,
0.1957+0.4219+0.3824=1,
...