Pytorch中torch.nn.Softmax的dim参数含义

涉及到多维tensor时,对softmax的参数dim总是很迷,下面用一个例子说明

import torch.nn as nn

m = nn.Softmax(dim=0)
n = nn.Softmax(dim=1)
k = nn.Softmax(dim=2)
input = torch.randn(2, 2, 3)
print(input)
print(m(input))
print(n(input))
print(k(input))

输出:
input

tensor([[[ 0.5450, -0.6264,  1.0446],
         [ 0.6324,  1.9069,  0.7158]],

        [[ 1.0092,  0.2421, -0.8928],
         [ 0.0344,  0.9723,  0.4328]]])

dim=0

tensor([[[0.3860, 0.2956, 0.8741],
         [0.6452, 0.7180, 0.5703]],

        [[0.6140, 0.7044, 0.1259],
         [0.3548, 0.2820, 0.4297]]])

dim=0时,在第0维上sum=1,即:
[0][0][0]+[1][0][0]=0.3860+0.6140=1
[0][0][1]+[1][0][1]=0.2956+0.7044=1
… …

dim=1

tensor([[[0.4782, 0.0736, 0.5815],
         [0.5218, 0.9264, 0.4185]],

        [[0.7261, 0.3251, 0.2099],
         [0.2739, 0.6749, 0.7901]]])

dim=1时,在第1维上sum=1,即:
[0][0][0]+[0][1][0]=0.4782+0.5218=1
[0][0][1]+[0][1][1]=0.0736+0.9264=1
… …

dim=2

tensor([[[0.3381, 0.1048, 0.5572],
         [0.1766, 0.6315, 0.1919]],

        [[0.6197, 0.2878, 0.0925],
         [0.1983, 0.5065, 0.2953]]])

dim=2时,在第2维上sum=1,即:
[0][0][0]+[0][0][1]+[0][0][2]=0.3381+0.1048+0.5572=1.0001(四舍五入问题)
[0][1][0]+[0][1][1]+[0][1][2]=0.1766+0.6315+0.1919=1
… …

用图表示223的张量如下:
在这里插入图片描述

  • 77
    点赞
  • 159
    收藏
    觉得还不错? 一键收藏
  • 7
    评论
评论 7
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值