File size: 2,622 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 | """RSC -- Representation Self-Challenging (Huang et al., ECCV 2020) -- DG ablation.
Each step, mute the features with the largest loss-gradient -- the dominant,
most-relied-upon directions, i.e. exactly the studio shortcut -- and force the
prediction from the remainder, so the model cannot lean on a single acquisition
cue. This module is the core masking primitive; the trainer computes the feature
gradient, builds the mask on a random subset of the batch, and re-forwards through
the head. Identity at eval (the trainer skips it).
"""
import torch
import torch.nn.functional as F
def rsc_feature_mask(grad: torch.Tensor, drop_frac: float = 1.0 / 3.0) -> torch.Tensor:
"""Return a 0/1 mask (shape of ``grad``) that zeros, per sample, the
``drop_frac`` fraction of features with the largest absolute gradient."""
_, c = grad.shape
k = int(round(drop_frac * c))
mask = torch.ones_like(grad)
if k <= 0:
return mask
_, idx = grad.abs().topk(k, dim=1)
return mask.scatter(1, idx, 0.0)
def rsc_step(feature_fn, head, imgs, labels,
drop_frac: float = 1.0 / 3.0, apply_frac: float = 0.5):
"""One Representation Self-Challenging step (Huang et al., ECCV 2020).
``feature_fn(imgs)`` gives pooled features ``[B, C]`` and ``head(features)`` the
logits ``[B, K]``. For a random ``apply_frac`` of the batch we mute, per sample,
the ``drop_frac`` fraction of channels carrying the largest loss-gradient, the
most-relied-upon directions, then predict from the remainder so the network
cannot lean on a single dominant (acquisition) cue. Returns
``(loss, logits_challenged)``. With ``drop_frac <= 0`` or ``apply_frac <= 0`` it
is plain cross-entropy, so the method reduces exactly to ERM. Use only in
training; at eval the caller does an ordinary forward.
"""
features = feature_fn(imgs)
if drop_frac <= 0 or apply_frac <= 0:
logits = head(features)
return F.cross_entropy(logits, labels), logits
loss_pre = F.cross_entropy(head(features), labels)
grad = torch.autograd.grad(loss_pre, features, retain_graph=True)[0]
mask = rsc_feature_mask(grad, drop_frac) # zeros top-k channels per row
if apply_frac < 1.0: # challenge a random subset only
b = features.size(0)
keep = torch.rand(b, device=features.device) >= apply_frac
mask = torch.where(keep.view(b, 1), torch.ones_like(mask), mask)
logits = head(features * mask)
return F.cross_entropy(logits, labels), logits
|