| 1 | from enum import Enum |
| 2 | |
| 3 | import torch |
| 4 | from torch import nn |
| 5 | |
| 6 | |
| 7 | class NormType(Enum): |
| 8 | """Normalization layer types: GROUP (GroupNorm) or PIXEL (per-location RMS norm).""" |
| 9 | |
| 10 | GROUP = "group" |
| 11 | PIXEL = "pixel" |
| 12 | |
| 13 | |
| 14 | class PixelNorm(nn.Module): |
| 15 | """ |
| 16 | Per-pixel (per-location) RMS normalization layer. |
| 17 | For each element along the chosen dimension, this layer normalizes the tensor |
| 18 | by the root-mean-square of its values across that dimension: |
| 19 | y = x / sqrt(mean(x^2, dim=dim, keepdim=True) + eps) |
| 20 | """ |
| 21 | |
| 22 | def __init__(self, dim: int = 1, eps: float = 1e-8) -> None: |
| 23 | """ |
| 24 | Args: |
| 25 | dim: Dimension along which to compute the RMS (typically channels). |
| 26 | eps: Small constant added for numerical stability. |
| 27 | """ |
| 28 | super().__init__() |
| 29 | self.dim = dim |
| 30 | self.eps = eps |
| 31 | |
| 32 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 33 | """ |
| 34 | Apply RMS normalization along the configured dimension. |
| 35 | """ |
| 36 | # Compute mean of squared values along `dim`, keep dimensions for broadcasting. |
| 37 | mean_sq = torch.mean(x**2, dim=self.dim, keepdim=True) |
| 38 | # Normalize by the root-mean-square (RMS). |
| 39 | rms = torch.sqrt(mean_sq + self.eps) |
| 40 | return x / rms |
| 41 | |
| 42 | |
| 43 | def build_normalization_layer( |
| 44 | in_channels: int, *, num_groups: int = 32, normtype: NormType = NormType.GROUP |
| 45 | ) -> nn.Module: |
| 46 | """ |
| 47 | Create a normalization layer based on the normalization type. |
| 48 | Args: |
| 49 | in_channels: Number of input channels |
| 50 | num_groups: Number of groups for group normalization |
| 51 | normtype: Type of normalization: "group" or "pixel" |
| 52 | Returns: |
| 53 | A normalization layer |
| 54 | """ |
| 55 | if normtype == NormType.GROUP: |
| 56 | return torch.nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True) |
| 57 | if normtype == NormType.PIXEL: |
| 58 | return PixelNorm(dim=1, eps=1e-6) |
| 59 | raise ValueError(f"Invalid normalization type: {normtype}") |
| 60 |