PyTorch torch.nn.Softmax 関数

PyTorch torch.nn リファレンスマニュアル

torch.nn.SoftmaxPyTorch の Softmax 活性化関数です。

入力を受け取り確率分布に変換します。すべての出力の合計は 1 になります。

関数定義

torch.nn.Softmax(dim=None)

パラメータ：

  • dim: Softmax を適用する次元

数式

Softmax(x_i) = exp(x_i) / sum(exp(x_j))

使用例

例1: 基本的な使い方

例

import torch
import torch.nn as nn

softmax = nn.Softmax(dim=1)

logits = torch.tensor([[2.0, 1.0, 0.1], [1.0, 2.0, 0.5]])
probs = softmax(logits)

print("Logits:", logits.tolist())
print("確率:", probs.tolist())
print("各行の合計:", probs.sum(dim=1).tolist())

例2: dim パラメータ

例

import torch
import torch.nn as nn

# 3D 入力
x = torch.randn(2, 3, 4)

# 異なる次元で softmax
print("dim=1:", nn.Softmax(dim=1)(x).sum(dim=1)[:1])
print("dim=2:", nn.Softmax(dim=2)(x).sum(dim=2)[:1])
print("dim=-1:", nn.Softmax(dim=-1)(x).sum(dim=-1)[:1])

例3: 分類出力

例

import torch
import torch.nn as nn

model = nn.Linear(784, 10)

# logits を出力
logits = model(torch.randn(4, 784))

# 確率に変換
probs = nn.Softmax(dim=1)(logits)

print("予測クラス:", probs.argmax(dim=1).tolist())
print("最大確率:", probs.max(dim=1).values.tolist())

使用シーン

  • 多クラス分類の出力: 確率分布
  • アテンション機構
  • 確率モデル

注意：Softmax 後の値はすべて正の数で、合計は 1 になります。

PyTorch torch.nn リファレンスマニュアル

その他の拡張