wehub-resource-sync
40 lines
1.0 KiB
Python
Executable File
40 lines
1.0 KiB
Python
Executable File
import numpy as np
|
|
from prml.nn.tensor.constant import Constant
|
|
from prml.nn.tensor.tensor import Tensor
|
|
from prml.nn.function import Function
|
|
|
|
|
|
class Softmax(Function):
|
|
|
|
def __init__(self, axis=-1):
|
|
if not isinstance(axis, int):
|
|
raise TypeError("axis must be int")
|
|
self.axis = axis
|
|
|
|
def _softmax(self, array):
|
|
y = array - np.max(array, self.axis, keepdims=True)
|
|
np.exp(y, out=y)
|
|
y /= y.sum(self.axis, keepdims=True)
|
|
return y
|
|
|
|
def forward(self, x):
|
|
x = self._convert2tensor(x)
|
|
self.x = x
|
|
self.output = self._softmax(x.value)
|
|
if isinstance(x, Constant):
|
|
return Constant(self.output)
|
|
return Tensor(self.output, function=self)
|
|
|
|
def backward(self, delta):
|
|
dx = self.output * delta
|
|
dx -= self.output * dx.sum(self.axis, keepdims=True)
|
|
self.x.backward(dx)
|
|
|
|
|
|
def softmax(x, axis=-1):
|
|
"""
|
|
softmax function along specified axis
|
|
y_k = exp(x_k) / sum_i(exp(x_i))
|
|
"""
|
|
return Softmax(axis=axis).forward(x)
|