Функция softmax

Функция softmax применяется к тензору и преобразует его элементы в вероятностное распределение. Первым параметром функция принимает тензор входных значений. Вторым параметром передается размерность dim, по которой выполняется нормализация. Результат функции - тензор той же формы, где значения лежат в диапазоне от 0 до 1, а их сумма по указанной размерности равна 1.

Синтаксис

torch.softmax(input, dim)

Пример

Давайте применим функцию softmax к одномерному тензору чисел:

import torch

t = torch.tensor([1.0, 2.0, 3.0])
res = torch.softmax(t, dim=0)
print(res)

Результат выполнения кода:

tensor([0.0900, 0.2447, 0.6652])

Сумма полученных вероятностей равна единице.

Пример

Применим функцию softmax к двумерному тензору по разным размерностям:

import torch

t = torch.tensor([
	[1.0, 2.0, 3.0],
	[4.0, 5.0, 6.0]
])

res_0 = torch.softmax(t, dim=0)
res_1 = torch.softmax(t, dim=1)

print("По строкам (dim=0):")
print(res_0)
print("\nПо столбцам (dim=1):")
print(res_1)

Результат выполнения кода:

"По строкам (dim=0):"
tensor([
	[0.0474, 0.0474, 0.0474],
	[0.9526, 0.9526, 0.9526]
])

"По столбцам (dim=1):"
tensor([
	[0.0900, 0.2447, 0.6652],
	[0.0900, 0.2447, 0.6652]
])

Пример

Функция softmax часто используется в нейросетях для многоклассовой классификации. Давайте применим ее к выходным данным модели:

import torch

torch.manual_seed(0)
logits = torch.randn(1, 4)
probabilities = torch.softmax(logits, dim=1)

print("Исходные логиты:")
print(logits)
print("\nВероятности классов:")
print(probabilities)
print("\nСумма вероятностей:")
print(probabilities.sum(dim=1))

Результат выполнения кода:

"Исходные логиты:"
tensor([[1.5410, -0.2934, -2.1788, 0.5684]])

"Вероятности классов:"
tensor([[0.5777, 0.0920, 0.0141, 0.3162]])

"Сумма вероятностей:"
tensor([1.0000])

Смотрите также

функцию sigmoid,
которая применяет сигмоидальное преобразование к элементам тензора
функцию log_softmax,
которая вычисляет логарифм от softmax для численной устойчивости
функцию log,
которая вычисляет натуральный логарифм элементов тензора
функцию exp,
которая вычисляет экспоненту элементов тензора