File size: 1,395 Bytes
1ed770c | 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 | """Channel-wise RMSNorm for NCHW tensors."""
from __future__ import annotations
import torch
from torch import Tensor, nn
class ChannelWiseRMSNorm(nn.Module):
"""Channel-wise RMSNorm with float32 reduction for numerical stability.
Normalizes across channels per spatial position. Supports optional
per-channel affine weight and bias.
"""
def __init__(self, channels: int, eps: float = 1e-6, affine: bool = True) -> None:
super().__init__()
self.channels: int = int(channels)
self._eps: float = float(eps)
if affine:
self.weight = nn.Parameter(torch.ones(self.channels))
self.bias = nn.Parameter(torch.zeros(self.channels))
else:
self.register_parameter("weight", None)
self.register_parameter("bias", None)
def forward(self, x: Tensor) -> Tensor:
if x.dim() < 2:
return x
# Float32 accumulation for stability
ms = torch.mean(torch.square(x), dim=1, keepdim=True, dtype=torch.float32)
inv_rms = torch.rsqrt(ms + self._eps)
y = x * inv_rms
if self.weight is not None:
shape = (1, -1) + (1,) * (x.dim() - 2)
y = y * self.weight.view(shape).to(dtype=y.dtype)
if self.bias is not None:
y = y + self.bias.view(shape).to(dtype=y.dtype)
return y.to(dtype=x.dtype)
|