File size: 2,392 Bytes
1ea7ba6 | 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 | """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
|