剖析 | torch.nn.functional.softmax维度详解

写代码，看代码都要心中有数，输入是什么，输出是什么，结果是如何计算出来的。

一维数据：

# -*- coding: utf-8 -*-
import torch
import numpy as np
import torch.nn.functional as F

a = [1, 2, 3]
a = torch.from_numpy(np.array(a, dtype=np.float))
b1 = F.softmax(a, dim=0)
b2 = F.softmax(a, dim=-1)
print(b1)
print(b2)

运行结果：

 dim=0,列和为1，列的和为1，由于只有一个维度，行在规模上是没有的，虽然眼睛看起来是1行3个数，但是实际上只有一个列维度是3，所以是列和为1，一共3列，那么计算出来的三个数的和应该是1，运行结果中可以看出，确实是，dim=-1和dim=0的结果一致，是循环所致。

二维数据：

# -*- coding: utf-8 -*-
import torch
import numpy as np
import torch.nn.functional as F

a = [[1, 2, 3],[4,5,6]]
a = torch.from_numpy(np.array(a, dtype=np.float))
b1 = F.softmax(a, dim=0)
b2 = F.softmax(a, dim=1)
b3 = F.softmax(a, dim=-1)
print(b1)
print(b2)
print(b3)

运行结果：

 

dim=-1和dim=1的结果一样，循环所致。 

三维数据：

# -*- coding: utf-8 -*-
import torch
import numpy as np
import torch.nn.functional as F

a = [[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [1, 2, 4]]]
a = torch.from_numpy(np.array(a, dtype=np.float))
b1 = F.softmax(a, dim=0)
b2 = F.softmax(a, dim=1)
b3 = F.softmax(a, dim=2)
b4 = F.softmax(a, dim=-1)
print(b1)
print(b2)
print(b3)
print(b4)

运行结果：

dim=-1和dim=2的结果一样，原因嘛，循环所致。