File size: 1,068 Bytes
38e555d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 | """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()
|