| """Supervised contrastive loss on existing fixed-crop foreground embeddings.""" |
| import torch |
| import torch.nn.functional as F |
|
|
| def pixel_contrast(features, target): |
| embeddings = features.float().permute(0,2,3,1).reshape(-1,features.shape[1]) |
| labels = target.reshape(-1) |
| selected=[] |
| for cls in range(1,19): |
| indices = (labels==cls).nonzero().flatten() |
| if len(indices): |
| selected.append(indices[torch.randperm(len(indices),device=indices.device)[:16]]) |
| if not selected: return features.sum()*0 |
| indices=torch.cat(selected) |
| embeddings=F.normalize(embeddings[indices],dim=1,eps=1e-4) |
| labels=labels[indices] |
| logits=embeddings@embeddings.T/.1 |
| eye=torch.eye(len(indices),device=features.device,dtype=torch.bool) |
| positives=(labels[:,None]==labels[None,:])&~eye |
| valid=positives.any(1) |
| if not valid.any(): return features.sum()*0 |
| log_probs=logits-logits.masked_fill(eye,float('-inf')).logsumexp(1,keepdim=True) |
| return -((log_probs*positives).sum(1)/positives.sum(1).clamp_min(1))[valid].mean() |
|
|
