Fonction PyTorch torch.nn.Softmax

PyTorch torch.nn 参考手册Manuel de référence torch.nn de PyTorch

torch.nn.Softmaxest la fonction d'activation Softmax dans PyTorch.

Il convertit l'entrée en une distribution de probabilité, avec toutes les sorties dont la somme vaut 1.

Définition de fonction

torch.nn.Softmax(dim=None)

Paramètres :

  • dim: La dimension selon laquelle Softmax est appliqué

Formule

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

Exemples d'utilisation

Exemple 1 : Utilisation de base

Exemple

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("Probabilités :", probs.tolist())
print("Sommes des lignes :", probs.sum(dim=1).tolist())

Exemple 2 : paramètre dim

Exemple

import torch
import torch.nn as nn

# entrée 3D
x = torch.randn(2, 3, 4)

# softmax selon différentes dimensions
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])

Exemple 3 : sortie de classification

Exemple

import torch
import torch.nn as nn

model = nn.Linear(784, 10)

# logits de sortie
logits = model(torch.randn(4, 784))

# Convertir en probabilités
probs = nn.Softmax(dim=1)(logits)

print("Classe prédite :", probs.argmax(dim=1).tolist())
print("Probabilité la plus élevée :", probs.max(dim=1).values.tolist())

Cas d'utilisation

  • Sortie de classification multi-classes: distribution de probabilité
  • Mécanisme d'attention
  • Modèles probabilistes

Remarque : Les valeurs après Softmax sont toutes positives et leur somme vaut 1.

PyTorch torch.nn 参考手册Manuel de référence torch.nn de PyTorch

Autres extensions