File size: 3,513 Bytes
ae73c7f | 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 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | """
Field normalization utilities compatible with The Well style (per-channel mean/std or RMS).
"""
from __future__ import annotations
import torch
from typing import Dict, Optional, Tuple
class FieldNormalizer:
def __init__(self, mode: str = "zscore", eps: float = 1e-6):
assert mode in ("zscore", "rms", "none")
self.mode = mode
self.eps = eps
self.mean: Optional[torch.Tensor] = None # (C,)
self.std: Optional[torch.Tensor] = None # (C,)
self.rms: Optional[torch.Tensor] = None # (C,)
self.fitted = False
@torch.no_grad()
def fit(self, data: torch.Tensor):
if isinstance(data, (list, tuple)):
data = torch.cat([d.reshape(-1, d.shape[-3], d.shape[-2], d.shape[-1]) for d in data], dim=0)
# Flatten to (..., C, H, W)
if data.dim() == 5: # (N, T, C, H, W)
data = data.permute(0, 1, 3, 4, 2).reshape(-1, data.shape[2])
elif data.dim() == 4: # (N, C, H, W)
data = data.permute(0, 2, 3, 1).reshape(-1, data.shape[1])
else:
raise ValueError(f"Unsupported shape {data.shape}")
self.mean = data.mean(dim=0)
self.std = data.std(dim=0).clamp_min(self.eps)
self.rms = torch.sqrt((data ** 2).mean(dim=0)).clamp_min(self.eps)
self.fitted = True
return self
def transform(self, x: torch.Tensor) -> torch.Tensor:
"""x: (..., C, H, W)"""
if self.mode == "none" or not self.fitted:
return x
actual_c = x.shape[-3]
expected_c = self.mean.shape[0]
if actual_c != expected_c:
raise ValueError(
f"FieldNormalizer was fit on {expected_c} channels but received "
f"data with {actual_c} channels. A normalizer fit on one domain "
f"cannot be reused for a domain with a different channel count "
f"-- fit a new FieldNormalizer for this domain instead of "
f"reusing the previous one. (This commonly happens in continual "
f"learning when successive domains have different channel "
f"counts; MultiScaleEncoder itself now handles varying channel "
f"counts, but normalization statistics are still per-domain.)"
)
# Broadcast stats to (C, 1, 1)
if self.mode == "zscore":
mean = self.mean.view(-1, 1, 1).to(x.device, x.dtype)
std = self.std.view(-1, 1, 1).to(x.device, x.dtype)
return (x - mean) / std
else: # rms
rms = self.rms.view(-1, 1, 1).to(x.device, x.dtype)
return x / rms
def inverse(self, x: torch.Tensor) -> torch.Tensor:
if self.mode == "none" or not self.fitted:
return x
if self.mode == "zscore":
mean = self.mean.view(-1, 1, 1).to(x.device, x.dtype)
std = self.std.view(-1, 1, 1).to(x.device, x.dtype)
return x * std + mean
else:
rms = self.rms.view(-1, 1, 1).to(x.device, x.dtype)
return x * rms
def state_dict(self) -> Dict:
return {
"mode": self.mode,
"mean": self.mean,
"std": self.std,
"rms": self.rms,
"fitted": self.fitted,
}
def load_state_dict(self, d: Dict):
self.mode = d["mode"]
self.mean = d["mean"]
self.std = d["std"]
self.rms = d["rms"]
self.fitted = d["fitted"]
|