PyTorch中的Softmax函数：用法和例子

在PyTorch中，torch.nn.Softmax()函数用于对输入张量进行softmax归一化操作。softmax函数可以将一组数值转换成概率分布的形式，常用于多分类问题中。
torch.nn.Softmax()函数的语法如下：

torch.nn.Softmax(dim, dtype=None, device=None, non_blocking=False)

参数说明：

• dim：指定进行softmax归一化的维度。可选值为0、1、2等，分别表示对输入张量的第0、1、2维度进行归一化。
• dtype：输出张量的数据类型，默认为输入张量的数据类型。
• device：输出张量所在的设备，默认为输入张量所在的设备。
• non_blocking：是否使用非阻塞模式进行计算，默认为False。
  下面通过几个例子来演示如何使用torch.nn.Softmax()函数：
  例1：对一维张量进行softmax归一化

  import torch
  import torch.nn.functional as F
  # 创建一个一维张量
  input_tensor = torch.tensor([2.0, 3.0, 4.0])
  # 对输入张量进行softmax归一化，dim=0表示对第0维度进行归一化
  output_tensor = F.softmax(input_tensor, dim=0)
  print(output_tensor)

  输出结果：

  tensor([0.090031, 0.244728, 0.665241])

  例2：对二维张量进行softmax归一化

  import torch
  import torch.nn.functional as F
  # 创建一个二维张量
  input_tensor = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
  # 对输入张量进行softmax归一化，dim=1表示对第1维度进行归一化
  output_tensor = F.softmax(input_tensor, dim=1)
  print(output_tensor)

  输出结果：

  tensor([[0.090031, 0.244728, 0.665241],
  [0.090031, 0.244728, 0.665241]])

  例3：对三维张量进行softmax归一化
  假设我们有一个三维张量，形状为(2, 3, 4)，需要进行softmax归一化操作。下面是一个示例代码：

  import torch
  import torch.nn.functional as F
  # 创建一个三维张量，形状为(2, 3, 4)
  input_tensor = torch.randn(2, 3, 4)
  # 对输入张量进行softmax归一化，dim=1表示对第1维度进行归一化
  output_tensor = F.softmax(input_tensor, dim=1)
  print(output_tensor)