nn.softmax

最近学习时遇到nn.softmax使用时的维度问题，记录一下。

nn.Softmax(dim=None)

• dim: 计算的维度

用softmax函数将N维输入进行归一化，归一化之后每个输出的Tensor范围在[0, 1]，并且归一化的那一维和为1
当输入数据是二维时，dim=0表示对列进行归一化，dim=1表示对行进行归一化

import torch.nn as nn
import torch

input = torch.Tensor([[1,2,3], [4,5,6], [7,8,9]])
m0 = nn.Softmax(dim=0)  # 对列
m1 = nn.Softmax(dim=1)  # 对行
output0 = m0(input)
output1 = m1(input)

print("input: ", input)
print("output0: ", output0)
print("output1: ", output1)

输出结果如下：

当输入数据是三维（更高维以此类推）

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

x= torch.Tensor([
    [[1,2,3],
     [4,5,6],
     [7,8,9]],
     [[1,2,3],
     [4,5,6],
     [7,8,9]],
])
m0 = nn.Softmax(dim=0)
m1 = nn.Softmax(dim=1)
m2 = nn.Softmax(dim=2)
m = nn.Softmax(dim=-1)
output0 = m0(x)
output1 = m1(x)
output2 = m2(x)
output = m(x)
print(x)
print("dim=0")
print(output0)
print("dim=1")
print(output1)
print("dim=2")
print(output2)
print("dim=-1")
print(output)

输出结果如下：

dim=2和dim=-1的结果一样

参考：https://www.zhihu.com/question/456213566?utm_id=0