Функция PyTorch torch.nn.Softmax

PyTorch torch.nn 参考手册Справочное руководство PyTorch torch.nn

torch.nn.Softmaxявляется функцией активации Softmax в PyTorch.

Она преобразует входные данные в распределение вероятностей, при этом сумма всех выходных значений равна 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.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 = 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 参考手册Справочное руководство PyTorch torch.nn

Другие расширения