torch.nn.functional的F.softmax()用法

F.softmax()就是softmax函数。

import torch
import torch.nn.functional as F
x = torch.rand(1,3,2)#1个3x2大小的矩阵
y = torch.rand(1,2,4)#1个2x4大小的矩阵
z = torch.bmm(x,y)
print('原始z矩阵：')
print(z)
print('按列方向上进行softmax：')
print(F.softmax(z, 1))
print('按行方向上进行softmax：')
print(F.softmax(z, 2))