File size: 2,378 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 | """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()
|