Función torch.nn.Softmax de PyTorch

PyTorch torch.nn 参考手册Referencia de torch.nn de PyTorch

torch.nn.SoftmaxEs la función de activación Softmax en PyTorch.

Convierte la entrada en una distribución de probabilidad, y la suma de todas las salidas es 1.

Definición de la función

torch.nn.Softmax(dim=None)

Parámetros:

  • dim: la dimensión sobre la que se realiza Softmax

Fórmula

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

Ejemplo de uso

Ejemplo 1: Uso básico

Ejemplo

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("Probabilidad:", probs.tolist())
print("Suma por fila:", probs.sum(dim=1).tolist())

Ejemplo 2: Parámetro dim

Ejemplo

import torch
import torch.nn as nn

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

# softmax según diferentes dimensiones
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])

Ejemplo 3: Salida de clasificación

Ejemplo

import torch
import torch.nn as nn

model = nn.Linear(784, 10)

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

# convertir a probabilidades
probs = nn.Softmax(dim=1)(logits)

print("Clase predicha:", probs.argmax(dim=1).tolist())
print("Probabilidad más alta:", probs.max(dim=1).values.tolist())

Escenarios de uso

  • Salida de clasificación múltiple: distribución de probabilidad
  • Mecanismo de atención
  • Modelos probabilísticos

Nota: Los valores después de Softmax son todos positivos y suman 1.

PyTorch torch.nn 参考手册Referencia de torch.nn de PyTorch

Otras extensiones