import math
from typing import Optional
import torch
from torch import Tensor, nn
from analogvnn.nn.activation.Activation import Activation
__all__ = ['PReLU', 'ReLU', 'LeakyReLU']
[docs]class PReLU(Activation):
"""Implements the parametric rectified linear unit (PReLU) activation function.
Attributes:
alpha (float): the slope of the negative part of the activation function.
_zero (Tensor): placeholder tensor of zero.
"""
[docs] __constants__ = ['alpha', '_zero']
def __init__(self, alpha: float):
"""Initialize the parametric rectified linear unit (PReLU) activation function.
Args:
alpha (float): the slope of the negative part of the activation function.
"""
super().__init__()
self.alpha = nn.Parameter(torch.tensor(alpha), requires_grad=False)
self._zero = nn.Parameter(torch.tensor(0), requires_grad=False)
[docs] def forward(self, x: Tensor) -> Tensor:
"""Forward pass of the parametric rectified linear unit (PReLU) activation function.
Args:
x (Tensor): the input tensor.
Returns:
Tensor: the output tensor.
"""
return torch.minimum(self._zero, x) * self.alpha + torch.maximum(self._zero, x)
[docs] def backward(self, grad_output: Optional[Tensor]) -> Optional[Tensor]:
"""Backward pass of the parametric rectified linear unit (PReLU) activation function.
Args:
grad_output (Optional[Tensor]): the gradient of the output tensor.
Returns:
Optional[Tensor]: the gradient of the input tensor.
"""
x = self.inputs
grad = (x < 0).type(torch.float) * self.alpha + (x >= 0).type(torch.float)
return grad_output * grad
@staticmethod
[docs] def initialise(tensor: Tensor) -> Tensor:
"""Initialisation of tensor using kaiming uniform, gain associated with PReLU activation function.
Args:
tensor (Tensor): the tensor to be initialized.
Returns:
Tensor: the initialized tensor.
"""
return nn.init.kaiming_uniform(tensor, a=math.sqrt(5), nonlinearity='leaky_relu')
@staticmethod
[docs] def initialise_(tensor: Tensor) -> Tensor:
"""In-place initialisation of tensor using kaiming uniform, gain associated with PReLU activation function.
Args:
tensor (Tensor): the tensor to be initialized.
Returns:
Tensor: the initialized tensor.
"""
return nn.init.kaiming_uniform_(tensor, a=math.sqrt(5), nonlinearity='leaky_relu')
[docs]class ReLU(PReLU):
"""Implements the rectified linear unit (ReLU) activation function.
Attributes:
alpha (float): 0
"""
def __init__(self):
"""Initialize the rectified linear unit (ReLU) activation function."""
super().__init__(alpha=0)
@staticmethod
[docs] def initialise(tensor: Tensor) -> Tensor:
"""Initialisation of tensor using kaiming uniform, gain associated with ReLU activation function.
Args:
tensor (Tensor): the tensor to be initialized.
Returns:
Tensor: the initialized tensor.
"""
return nn.init.kaiming_uniform(tensor, a=math.sqrt(5), nonlinearity='relu')
@staticmethod
[docs] def initialise_(tensor: Tensor) -> Tensor:
"""In-place initialisation of tensor using kaiming uniform, gain associated with ReLU activation function.
Args:
tensor (Tensor): the tensor to be initialized.
Returns:
Tensor: the initialized tensor.
"""
return nn.init.kaiming_uniform_(tensor, a=math.sqrt(5), nonlinearity='relu')
[docs]class LeakyReLU(PReLU):
"""Implements the leaky rectified linear unit (LeakyReLU) activation function.
Attributes:
alpha (float): 0.01
"""
def __init__(self):
"""Initialize the leaky rectified linear unit (LeakyReLU) activation function."""
super().__init__(alpha=0.01)