Spaces:
Running on Zero
Running on Zero
File size: 1,263 Bytes
0122a25 | 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 | """Utility functions for layer ops."""
from __future__ import annotations
from torch import nn
def build_activation_layer(
activation: str, inplace: bool = False
) -> nn.Module:
"""Build activation layer.
Args:
activation (str): Activation layer type.
inplace (bool, optional): If to set inplace. Defaults to False. It will
be ignored if the activation layer is not inplace.
"""
activation_layer = getattr(nn, activation)
if activation_layer in {nn.Tanh, nn.PReLU, nn.Sigmoid, nn.GELU}:
return activation_layer()
return activation_layer(inplace=inplace)
def build_norm_layer(
norm: str, out_channels: int, num_groups: int | None = None
) -> nn.Module:
"""Build normalization layer.
Args:
norm (str): Normalization layer type.
out_channels (int): Number of output channels.
num_groups (int | None, optional): Number of groups for GroupNorm.
Defaults to None.
"""
norm_layer = getattr(nn, norm)
if norm_layer == nn.GroupNorm:
assert (
num_groups is not None
), "num_groups must be specified when using Group Norm"
return norm_layer(num_groups, out_channels)
return norm_layer(out_channels)
|