ml4gw.nn.norm

Classes

`GroupNorm1D`(num_channels[, num_groups, eps])	Custom implementation of GroupNorm which is faster than the out-of-the-box PyTorch version at inference time.
`GroupNorm1DGetter`([groups])	Utility for making a NormLayer Callable that maps from an integer number of channels to a torch Module.
`GroupNorm2D`(num_channels[, num_groups, eps, ...])	torch.nn.GroupNorm with the number of channels as its first argument, like GroupNorm1D.
`GroupNorm2DGetter`([groups])	Utility for making a NormLayer Callable that maps from an integer number of channels to a torch Module.

class ml4gw.nn.norm.GroupNorm1D(num_channels, num_groups=None, eps=1e-05)

Bases: Module

Custom implementation of GroupNorm which is faster than the out-of-the-box PyTorch version at inference time.

num_groups defaults to one group per channel, and is capped at num_channels, so a single value can be used for every layer of a network whose narrowest layer has fewer channels than that.

Parameters:

num_channels (int)
num_groups (int | None)
eps (float)

forward(x)

Return type:: Float[Tensor, 'batch channel length']
Parameters:: x (Float[Tensor, 'batch channel length'])

class ml4gw.nn.norm.GroupNorm1DGetter(groups=None)

Bases: object

Utility for making a NormLayer Callable that maps from an integer number of channels to a torch Module.

Deprecated since version Use: GroupNorm1D directly, e.g. functools.partial(GroupNorm1D, num_groups=groups) in code, or GroupNorm1D as the class path with num_groups in a jsonargparse config. jsonargparse>=4.52 rejects this class for a NormLayer.

Parameters:: groups (int | None)

class ml4gw.nn.norm.GroupNorm2D(num_channels, num_groups=None, eps=1e-05, affine=True)

Bases: GroupNorm

torch.nn.GroupNorm with the number of channels as its first argument, like GroupNorm1D.

This lets a NormLayer be specified with jsonargparse by pointing at this class and setting only num_groups: jsonargparse passes the channel count as the first positional argument when the layer is built. As with GroupNorm1D, num_groups defaults to one group per channel and is capped at num_channels.

Parameters:

num_channels (int)
num_groups (int)
eps (float)
affine (bool)

class ml4gw.nn.norm.GroupNorm2DGetter(groups=None)

Bases: object

Utility for making a NormLayer Callable that maps from an integer number of channels to a torch Module.

Deprecated since version Use: GroupNorm2D directly, e.g. functools.partial(GroupNorm2D, num_groups=groups) in code, or GroupNorm2D as the class path with num_groups in a jsonargparse config. jsonargparse>=4.52 rejects this class for a NormLayer.

Parameters:: groups (int | None)