| 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 | |