fastreid.layers.batch_norm Namespace Reference

Classes
class	BatchNorm
class	FrozenBatchNorm
class	GhostBatchNorm
class	IBN
class	SyncBatchNorm

Functions
	get_norm (norm, out_channels, **kwargs)

Detailed Description

@author:  liaoxingyu
@contact: sherlockliao01@gmail.com

Function Documentation

get_norm()

fastreid.layers.batch_norm.get_norm	(		norm,
			out_channels,
		**	kwargs )

Args:
    norm (str or callable): either one of BN, GhostBN, FrozenBN, GN or SyncBN;
        or a callable that thakes a channel number and returns
        the normalization layer as a nn.Module
    out_channels: number of channels for normalization layer

Returns:
    nn.Module or None: the normalization layer

Definition at line 187 of file batch_norm.py.

def get_norm(norm, out_channels, **kwargs):
    """
    Args:
        norm (str or callable): either one of BN, GhostBN, FrozenBN, GN or SyncBN;
            or a callable that thakes a channel number and returns
            the normalization layer as a nn.Module
        out_channels: number of channels for normalization layer
 
    Returns:
        nn.Module or None: the normalization layer
    """
    if isinstance(norm, str):
        if len(norm) == 0:
            return None
        norm = {
            "BN": BatchNorm,
            "GhostBN": GhostBatchNorm,
            "FrozenBN": FrozenBatchNorm,
            "GN": lambda channels, **args: nn.GroupNorm(32, channels),
            "syncBN": SyncBatchNorm,
        }[norm]
    return norm(out_channels, **kwargs)