Contrastive Learning대조 학습

같은 이미지의 두 증강 → 임베딩 공간에서 가깝게
다른 이미지들     → 임베딩 공간에서 멀게

                양성 쌍
  [이미지A] ──→ [증강1] ──→ 인코더 ──→ z1 ←──┐ 유사
              [증강2] ──→ 인코더 ──→ z2 ──→┘
  [이미지B] ──→ [증강3] ──→ 인코더 ──→ z3   (음성 쌍과 멀게)

import torch
import torch.nn as nn
import torch.nn.functional as F

class SimCLR(nn.Module):
    def __init__(self, encoder, projection_dim=128):
        super().__init__()
        self.encoder = encoder
        # 투영 헤드: 표현 → 대조 학습 공간
        h_dim = encoder.fc.in_features
        self.projector = nn.Sequential(
            nn.Linear(h_dim, h_dim), nn.ReLU(),
            nn.Linear(h_dim, projection_dim)
        )
        self.encoder.fc = nn.Identity()

    def forward(self, x):
        h = self.encoder(x)
        z = self.projector(h)
        return F.normalize(z, dim=1)

def nt_xent_loss(z1, z2, temperature=0.5):
    """
    NT-Xent (Normalized Temperature-scaled Cross Entropy)
    배치 내 다른 샘플들을 음성 쌍으로 사용
    """
    N = z1.shape[0]
    z = torch.cat([z1, z2], dim=0)  # (2N, D)
    sim = torch.mm(z, z.T) / temperature  # (2N, 2N)

    # 자기 자신 제외 마스크
    mask = torch.eye(2*N, dtype=torch.bool, device=z.device)
    sim.masked_fill_(mask, float('-inf'))

    # 양성 쌍 인덱스
    labels = torch.cat([torch.arange(N, 2*N), torch.arange(N)]).to(z.device)
    return F.cross_entropy(sim, labels)

# 학습 루프
def train_simclr(model, loader, optimizer, epochs=100):
    augment = get_augmentation()
    for epoch in range(epochs):
        for images, _ in loader:
            x1 = augment(images).cuda()
            x2 = augment(images).cuda()
            z1 = model(x1)
            z2 = model(x2)
            loss = nt_xent_loss(z1, z2)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

def clip_loss(image_emb, text_emb, temperature=0.07):
    """
    이미지-텍스트 쌍을 정렬하는 대조 손실
    (배치 내 (i,i) 쌍이 양성)
    """
    image_emb = F.normalize(image_emb, dim=-1)
    text_emb  = F.normalize(text_emb, dim=-1)
    logits = (image_emb @ text_emb.T) / temperature
    N = logits.shape[0]
    labels = torch.arange(N, device=logits.device)
    loss_i = F.cross_entropy(logits, labels)
    loss_t = F.cross_entropy(logits.T, labels)
    return (loss_i + loss_t) / 2