Pytorch torch.nn.Softmax的dim参数

比如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,

...

  • 2
    点赞
  • 6
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值