| 28 | |
| 29 | |
| 30 | def get_activation_function(name: str, functional: bool = False, num: int = 1): |
| 31 | name = name.lower().strip() |
| 32 | |
| 33 | def get_functional(s: str) -> Optional[Callable]: |
| 34 | return {"softmax": F.softmax, "relu": F.relu, "tanh": torch.tanh, "sigmoid": torch.sigmoid, |
| 35 | "identity": nn.Identity(), |
| 36 | None: None, 'swish': F.silu, 'silu': F.silu, 'elu': F.elu, 'gelu': F.gelu, 'prelu': nn.PReLU(), |
| 37 | }[s] |
| 38 | |
| 39 | def get_nn(s: str) -> Optional[Callable]: |
| 40 | return {"softmax": nn.Softmax(dim=1), "relu": nn.ReLU(), "tanh": nn.Tanh(), "sigmoid": nn.Sigmoid(), |
| 41 | "identity": nn.Identity(), 'silu': nn.SiLU(), 'elu': nn.ELU(), 'prelu': nn.PReLU(), |
| 42 | 'swish': nn.SiLU(), 'gelu': nn.GELU(), |
| 43 | }[s] |
| 44 | |
| 45 | if num == 1: |
| 46 | return get_functional(name) if functional else get_nn(name) |
| 47 | else: |
| 48 | return [get_nn(name) for _ in range(num)] |
| 49 | |
| 50 | |
| 51 | def get_normalization_layer(name, dims, num_groups=None, *args, **kwargs): |