Files
PINA/pina/adaptive_functions/adaptive_exp.py
2024-02-09 15:11:51 +01:00

54 lines
1.6 KiB
Python

import torch
from torch.nn.parameter import Parameter
class AdaptiveExp(torch.nn.Module):
"""
Implementation of soft exponential activation.
Shape:
- Input: (N, *) where * means, any number of additional
dimensions
- Output: (N, *), same shape as the input
Parameters:
- alpha - trainable parameter
References:
- See related paper:
https://arxiv.org/pdf/1602.01321.pdf
Examples:
>>> a1 = soft_exponential(256)
>>> x = torch.randn(256)
>>> x = a1(x)
"""
def __init__(self):
"""
Initialization.
INPUT:
- in_features: shape of the input
- aplha: trainable parameter
aplha is initialized with zero value by default
"""
super(AdaptiveExp, self).__init__()
self.scale = Parameter(
torch.normal(torch.tensor(1.0), torch.tensor(0.1))
) # create a tensor out of alpha
self.scale.requiresGrad = True # set requiresGrad to true!
self.alpha = Parameter(
torch.normal(torch.tensor(1.0), torch.tensor(0.1))
) # create a tensor out of alpha
self.alpha.requiresGrad = True # set requiresGrad to true!
self.translate = Parameter(
torch.normal(torch.tensor(0.0), torch.tensor(0.1))
) # create a tensor out of alpha
self.translate.requiresGrad = True # set requiresGrad to true!
def forward(self, x):
"""
Forward pass of the function.
Applies the function to the input elementwise.
"""
return self.scale * (x + self.translate)