Pytorch tf.nn.functional.softmax(x,dim = -1)对参数dim的理解

最新推荐文章于 2024-10-28 11:12:22 发布

转载最新推荐文章于 2024-10-28 11:12:22 发布 · 2.7k 阅读

49 ·

CC 4.0 BY-SA版权

原文链接：https://blog.youkuaiyun.com/Will_Ye/article/details/104994504

文章标签：

#深度学习 #python #算法

Pytorch 专栏收录该内容

30 篇文章

订阅专栏

本文详细解析了PyTorch中Softmax函数的使用方法，通过实例演示了不同维度设置下Softmax运算的过程及结果，帮助读者深入理解该函数的工作原理。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

torch.nn.functional.Softmax(input,dim=None)
tf.nn.functional.softmax(x,dim)中的参数dim是指维度的意思，设置这个参数时会遇到0,1，2，-1等情况。

在这里插入图片描述

一般会有设置成dim=0,1,2,-1的情况。
准备工作：先随机生成一个（2,2,3）的矩阵，两个维度的（2,3）矩阵。

import torch
import torch.nn.functional as F
input = torch.randn(2,2,3))
print(input)

输出为：

tensor([[[-3.9332,  0.7909,  0.8927],
         [-1.7991,  0.2505,  0.7695]],

        [[ 0.1946,  0.1878,  1.2713],
         [ 0.9536,  1.0525, -0.7081]]])

dim=0是对每一维度相同位置的数值进行softmax运算

例如：

m = F.softmax(input,dim=0)
print(m)

输出为：

tensor([[[0.0159, 0.6464, 0.4065],
         [0.0599, 0.3096, 0.8142]],

        [[0.9841, 0.3536, 0.5935],
         [0.9401, 0.6904, 0.1858]]])

从输出结果可以看出，一共2个维度
每一个维度（2,3）对应的数值相加为1，例如：

[0][0][0]+[1][0][0]=0.0159+0.9841=1
[0][0][1]+[1][0][1]=0.6464+0.3536=1

dim=1是对某一维度的列进行softmax运算

m = F.softmax(input,dim=1)
print(m)

输出为：

tensor([[[0.1058, 0.6319, 0.5308],
         [0.8942, 0.3681, 0.4692]],

        [[0.3189, 0.2964, 0.8786],
         [0.6811, 0.7036, 0.1214]]])

从输出结果可以看出，计算的是某一维度中的列，例如：
在第1维度中：
[0][0][0]+[0][1][0]=0.1058+0.8942=1
在第2维度中：
[1][0][0]+[1][1][0]=0.3189+0.6811=1

dim=2是对某一维度的行进行softmax运算

m = F.softmax(input,dim=2)
print(m)

输出为：

tensor([[[0.0042, 0.4726, 0.5232],
         [0.0458, 0.3560, 0.5982]],

        [[0.2029, 0.2015, 0.5955],
         [0.4360, 0.4813, 0.0828]]])

从输出结果可以看出，计算的是某一维度中的行，例如：
在第1维度中：
[0][0][0]+[0][0][1]+[0][0][2]=0.0042+0.4726+0.5232=1
在第2维度中：
[1][0][0]+[1][0][1]+[1][0][2]=0.2029+0.2015+0.5955=0.9999=1（是等于1的，只是后面还有很多位小数省略了）