Класс Softmax

Класс Softmax из модуля torch.nn применяет функцию softmax к входному тензору вдоль заданной размерности. Это преобразование часто используется в задачах многоклассовой классификации для получения вероятностей принадлежности к классам. Первый параметр dim указывает ось, по которой производится нормализация. Второй параметр dtype позволяет задать тип данных выходного тензора.

Синтаксис

torch.nn.Softmax(dim, dtype=None)

Пример

Создадим слой softmax и применим его к тензору логитов:

import torch
import torch.nn as nn

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

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

tensor([0.0900, 0.2447, 0.6652])

Пример

Применим softmax к двумерному тензору по строкам:

import torch
import torch.nn as nn

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

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

tensor([
	[0.0900, 0.2447, 0.6652],
	[0.0900, 0.2447, 0.6652],
])

Пример

Используем параметр dtype для указания типа данных:

import torch
import torch.nn as nn

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

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

torch.float64

Пример

Применим softmax в составе последовательной модели:

import torch
import torch.nn as nn

model = nn.Sequential(
	nn.Linear(4, 3),
	nn.Softmax(dim=1)
)
t = torch.randn(2, 4)
res = model(t)
print(res)

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

tensor([
	[0.4246, 0.0533, 0.5221],
	[0.7972, 0.0908, 0.1120],
], grad_fn=<SoftmaxBackward0>)

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

функцию LogSoftmax,
которая применяет логарифм от softmax
функцию Softmin,
которая применяет softmin к входному тензору
функцию Sigmoid,
которая применяет сигмоидальную функцию активации
функцию ReLU,
которая применяет функцию активации ReLU