File size: 1,656 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 | """ARC-V training objective (ARCV_METHOD_DESIGN.md §4.3).
Combines the loss-side ARC-V mechanisms on top of any model that maps images to
logits: the Fourier amplitude-mix (M2) augmentation plus a FACT-style
phase-consistency co-teaching term. MixStyle (M1) lives inside the backbone via
:func:`src.dg.mixstyle.register_mixstyle_hooks` (architectural, no loss term); the
cross-source SupCon (M3) is a LOSO-only add-on for Regime B.
With ``p_f=0`` and ``gamma_pc=0`` this reduces *exactly* to ERM cross-entropy, so a
single code path covers both the ERM baseline and the ARC-V method (toggled by
config) -- the honest apples-to-apples comparison the design demands.
``forward`` is a callable mapping an image batch to class logits, e.g.
``model.forward_image``.
"""
import torch.nn.functional as F
from src.dg.fourier import FourierAmplitudeMix, phase_consistency_loss
class ARCVObjective:
def __init__(self, p_f: float = 0.5, eta: float = 1.0,
eta_weak: float = 0.3, gamma_pc: float = 1.0):
self.strong = FourierAmplitudeMix(p=p_f, eta=eta)
self.weak = FourierAmplitudeMix(p=p_f, eta=eta_weak)
self.gamma_pc = float(gamma_pc)
def __call__(self, forward, imgs, labels, training: bool = True):
self.strong.train(training)
self.weak.train(training)
logits = forward(self.strong(imgs))
loss = F.cross_entropy(logits, labels)
if training and self.gamma_pc > 0.0:
logits_weak = forward(self.weak(imgs))
loss = loss + self.gamma_pc * phase_consistency_loss(logits, logits_weak)
return loss, logits
|