剖析 | 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)

运行结果:

 剖析 | torch.nn.functional.softmax维度详解_第1张图片

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)

运行结果:

剖析 | torch.nn.functional.softmax维度详解_第2张图片

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

你可能感兴趣的:(pytorch)