| """Fourier amplitude-mix with phase preservation (FACT; Xu et al., CVPR 2021). | |
| Acquisition shift (white background, warm cast, JPEG, gamma) concentrates in the | |
| *amplitude* spectrum; granule shape/micro-texture lives in *phase*. During | |
| training, with probability ``p``, mix each image's amplitude spectrum with a | |
| shuffled peer's while KEEPING its own phase -- perturbing acquisition style | |
| without touching semantics. Identity at eval. ARC-V's M2 mechanism; the | |
| phase-stable output is also what the (previously dormant) frequency branch | |
| consumes. | |
| Pair with :func:`phase_consistency_loss` -- a FACT-style co-teaching term that | |
| keeps the prediction on a strongly amplitude-mixed image consistent with a | |
| weakly-mixed copy (they share phase, hence semantics). | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| class FourierAmplitudeMix(nn.Module): | |
| def __init__(self, p: float = 0.5, eta: float = 1.0): | |
| super().__init__() | |
| self.p = float(p) | |
| self.eta = float(eta) | |
| 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) | |
| f = torch.fft.rfft2(x) | |
| amp, phase = f.abs(), f.angle() | |
| perm = torch.randperm(b, device=x.device) | |
| lam = torch.empty(b, 1, 1, 1, device=x.device).uniform_(0.0, self.eta) | |
| amp_mix = lam * amp + (1.0 - lam) * amp[perm] | |
| out = torch.fft.irfft2(torch.polar(amp_mix, phase), s=x.shape[-2:]) | |
| return out.to(dtype=x.dtype) | |
| def extra_repr(self) -> str: | |
| return f"p={self.p}, eta={self.eta}" | |
| def phase_consistency_loss(logits_strong: torch.Tensor, | |
| logits_weak: torch.Tensor) -> torch.Tensor: | |
| """Symmetric Jensen-Shannon divergence between the class posteriors of a | |
| strongly amplitude-mixed image and a weakly-mixed copy. Both share this image's | |
| phase (semantics), so their predictions should agree (FACT co-teaching).""" | |
| p = F.softmax(logits_strong, dim=1) | |
| q = F.softmax(logits_weak, dim=1) | |
| m = (0.5 * (p + q)).clamp_min(1e-8) | |
| kl_pm = (p * (p.clamp_min(1e-8).log() - m.log())).sum(dim=1) | |
| kl_qm = (q * (q.clamp_min(1e-8).log() - m.log())).sum(dim=1) | |
| return (0.5 * (kl_pm + kl_qm)).mean() | |