[人工智能]Pytorch小记-torch.nn.Softmax(dim=1)如何理解？

在进行多分类时，时常看到Softmax分类函数，理论基础可看文章PyTorch深度学习实践概论笔记9-SoftMax分类器。应用时经常看到代码torch.nn.Softmax(dim=1)，如何理解呢？

1.文档

pytorch不会，查看官方文档：Softmax — PyTorch 1.10 documentation。

torch.nn.Softmax 是 Module 类，在实例化类后会初始化运算所需要的参数。这些参数会在 forward 和 backward 之后根据 loss 进行更新，通常存放在定义模型的 __init__() 中。在 Module 类里的 __call__ 实现了 forward() 函数的调用。

定义

将Softmax函数应用于一个n维输入张量，对其进行缩放，使n维输出张量的元素位于[0,1]范围内，总和为1。

公式

参数

dim (int) - Softmax将被计算的维度(因此沿dim的每个切片和为1)。

请注意

此模块不直接与NLLLoss一起工作，它期望在Softmax和它自己之间计算Log。

可以使用LogSoftmax代替(它更快并且有更好的数值属性)。

2.例子

看个栗子：

import torch
a = torch.randn(2,3,4)
print(a)

a是维度为（2,3,4）的tensor，输出为：

tensor([[[-0.5947,  0.9496, -1.9366,  0.0580],
         [ 0.8346,  2.2958,  1.4638, -0.0390],
         [ 0.1379,  1.2085,  1.2671, -1.2156]],

        [[-0.1980, -0.8921, -1.1665, -1.2323],
         [-0.8220, -1.9701,  1.7183, -0.2698],
         [ 0.1977, -0.1474, -1.5431, -2.2217]]])

当dim=0时，指的是在维度0上的元素相加等于1。

torch.nn.Softmax(dim=0)(a)

当dim=1时，指的是在维度1上的元素相加等于1。

torch.nn.Softmax(dim=1)(a)

当dim=2时，指的是在维度2上的元素相加等于1。

torch.nn.Softmax(dim=2)(a)

说明：记录学习笔记，如果错误欢迎指正！写文章不易，转载请联系我。