File size: 1,266 Bytes
338c3e4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 | from torch import nn
from torch import Tensor
def get_act_fn(name: str, norm: bool = False) -> nn.Module:
if name == "relu":
fn = nn.ReLU()
elif name == "tanh":
fn = nn.Tanh()
elif name == "gelu":
fn = nn.GELU()
elif name == "swish":
fn = nn.SiLU()
else:
raise ValueError(f"Unknown activation function: {name}")
if norm:
fn = NormAct(fn)
return fn
class NormAct(nn.Module):
"""
Normalized Activation Function.
A wrapper around any activation function that normalizes the input
before applying the activation function, and then transforms the
output back to the original scale.
"""
def __init__(self, act_fn: nn.Module):
super().__init__()
self.act_fn = act_fn
def forward(self, x: Tensor) -> Tensor:
'''
x: (b, h, w)
'''
num_dims = len(x.shape)
dims = tuple(range(1, num_dims))
# find the mean and std of each example in the batch
mean = x.mean(dim=dims, keepdim=True)
std = x.std(dim=dims, keepdim=True)
# normalize
x = (x - mean) / std
x = self.act_fn(x)
# Transform back to the original scale
x = x * std + mean
return x
|