| """MixStyle (Zhou et al., ICLR 2021) -- feature-statistic randomization. | |
| During training, with probability ``p``, replace each sample's per-channel | |
| statistics ``(mu, sigma)`` with a random convex combination of its own and a | |
| shuffled peer's, then re-apply -- synthesizing pseudo-acquisition-domains the | |
| single training source lacks. Identity at eval. This is ARC-V's M1 mechanism and | |
| also a standalone DG baseline; insert as a hook after early backbone stages. | |
| The operation rewrites only per-channel affine statistics, so the | |
| instance-normalized *content* is preserved by construction. | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| class MixStyle(nn.Module): | |
| def __init__(self, p: float = 0.5, alpha: float = 0.1, eps: float = 1e-6): | |
| super().__init__() | |
| self.p = float(p) | |
| self.alpha = float(alpha) | |
| self.eps = float(eps) | |
| self._beta = torch.distributions.Beta(self.alpha, self.alpha) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| if not self.training or self.p <= 0.0: | |
| return x | |
| if torch.rand(1).item() > self.p: | |
| return x | |
| b = x.size(0) | |
| mu = x.mean(dim=(2, 3), keepdim=True) | |
| sig = (x.var(dim=(2, 3), keepdim=True, unbiased=False) + self.eps).sqrt() | |
| x_norm = (x - mu) / sig | |
| lam = self._beta.sample((b, 1, 1, 1)).to(device=x.device, dtype=x.dtype) | |
| perm = torch.randperm(b, device=x.device) | |
| mu_mix = lam * mu + (1.0 - lam) * mu[perm] | |
| sig_mix = lam * sig + (1.0 - lam) * sig[perm] | |
| return x_norm * sig_mix + mu_mix | |
| def extra_repr(self) -> str: | |
| return f"p={self.p}, alpha={self.alpha}" | |
| def register_mixstyle_hooks(modules, p: float = 0.5, alpha: float = 0.1): | |
| """Attach a MixStyle forward-hook to each module in ``modules`` (e.g. the first | |
| few backbone stages) so their outputs are style-randomized during training and | |
| passed through unchanged at eval. The hook syncs the MixStyle to the hooked | |
| module's train/eval mode and replaces its output. Returns the hook handles.""" | |
| handles = [] | |
| for m in modules: | |
| ms = MixStyle(p=p, alpha=alpha) | |
| def _hook(module, inputs, output, _ms=ms): | |
| _ms.train(module.training) | |
| return _ms(output) | |
| handles.append(m.register_forward_hook(_hook)) | |
| return handles | |