artist_style_dinov3/losses.py

from __future__ import annotations

import math
from typing import Dict

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


def _off_diagonal_mask(size: int, device: torch.device) -> torch.Tensor:
    mask = torch.ones(size, size, dtype=torch.bool, device=device)
    mask.fill_diagonal_(False)
    return mask


def supervised_contrastive_loss(embeddings: torch.Tensor, labels: torch.Tensor, temperature: float) -> torch.Tensor:
    if embeddings.size(0) <= 1:
        return embeddings.new_zeros(())

    normalized = F.normalize(embeddings.float(), dim=-1)
    logits = torch.matmul(normalized, normalized.t()) / temperature
    logits_mask = ~torch.eye(logits.size(0), dtype=torch.bool, device=logits.device)
    logits = logits - logits.max(dim=1, keepdim=True).values.detach()

    label_mask = labels.unsqueeze(0).eq(labels.unsqueeze(1)) & logits_mask
    positives_per_row = label_mask.sum(dim=1)
    valid_rows = positives_per_row > 0
    if not valid_rows.any():
        return embeddings.new_zeros(())

    exp_logits = torch.exp(logits) * logits_mask.to(dtype=logits.dtype)
    log_prob = logits - torch.log(exp_logits.sum(dim=1, keepdim=True).clamp_min(1e-12))
    positive_log_prob = log_prob.masked_fill(~label_mask, 0.0)
    loss = -positive_log_prob.sum(dim=1) / positives_per_row.clamp_min(1)
    return loss[valid_rows].mean().to(dtype=embeddings.dtype)


class MultiPrototypeLoss(nn.Module):
    def __init__(
        self,
        classification_weight: float = 1.0,
        arcface_scale: float = 30.0,
        arcface_margin: float = 0.30,
        compactness_weight: float = 0.25,
        diversity_weight: float = 0.05,
        usage_weight: float = 0.05,
        branch_orthogonality_weight: float = 0.05,
        branch_balance_weight: float = 0.05,
        view_balance_weight: float = 0.05,
        supcon_weight: float = 0.1,
        label_smoothing: float = 0.0,
        assignment_temperature: float = 0.1,
        supcon_temperature: float = 0.1,
        prototype_margin: float = 0.15,
    ) -> None:
        super().__init__()
        self.classification_weight = classification_weight
        self.arcface_scale = float(arcface_scale)
        self.arcface_margin = float(arcface_margin)
        self.compactness_weight = compactness_weight
        self.diversity_weight = diversity_weight
        self.usage_weight = usage_weight
        self.branch_orthogonality_weight = branch_orthogonality_weight
        self.branch_balance_weight = branch_balance_weight
        self.view_balance_weight = view_balance_weight
        self.supcon_weight = supcon_weight
        self.label_smoothing = label_smoothing
        self.assignment_temperature = assignment_temperature
        self.supcon_temperature = supcon_temperature
        self.prototype_margin = prototype_margin

    def _arcface_loss(self, cosine: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
        cosine = cosine.float().clamp(-1.0 + 1e-6, 1.0 - 1e-6)
        sine = torch.sqrt((1.0 - cosine.pow(2)).clamp_min(1e-6))
        cos_m = math.cos(self.arcface_margin)
        sin_m = math.sin(self.arcface_margin)
        th = math.cos(math.pi - self.arcface_margin)
        mm = math.sin(math.pi - self.arcface_margin) * self.arcface_margin

        phi = cosine * cos_m - sine * sin_m
        phi = torch.where(cosine > th, phi, cosine - mm)

        one_hot = torch.zeros_like(cosine)
        one_hot.scatter_(1, labels.view(-1, 1), 1.0)
        logits = (one_hot * phi + (1.0 - one_hot) * cosine) * self.arcface_scale
        return F.cross_entropy(logits, labels, label_smoothing=self.label_smoothing)

    def _compactness_loss(self, true_sims: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        assignments = torch.softmax(true_sims / self.assignment_temperature, dim=-1)
        compactness = (1.0 - (assignments * true_sims).sum(dim=-1)).mean()
        return compactness, assignments

    def _prototype_diversity_loss(self, normalized_prototypes: torch.Tensor) -> torch.Tensor:
        pairwise = torch.matmul(normalized_prototypes, normalized_prototypes.transpose(-1, -2))
        mask = _off_diagonal_mask(pairwise.size(-1), pairwise.device)
        off_diag = pairwise[:, mask].view(pairwise.size(0), -1)
        return F.relu(off_diag - self.prototype_margin).pow(2).mean()

    def _prototype_usage_loss(self, assignments: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
        losses = []
        target = torch.full(
            (assignments.size(-1),),
            1.0 / assignments.size(-1),
            device=assignments.device,
            dtype=assignments.dtype,
        )
        for label in labels.unique():
            mask = labels == label
            if mask.sum() < 2:
                continue
            mean_assignment = assignments[mask].mean(dim=0)
            losses.append((mean_assignment - target).pow(2).mean())
        if not losses:
            return assignments.new_zeros(())
        return torch.stack(losses).mean()

    def _branch_orthogonality_loss(self, branch_embeddings: torch.Tensor) -> torch.Tensor:
        normalized = F.normalize(branch_embeddings, dim=-1)
        gram = torch.matmul(normalized, normalized.transpose(-1, -2))
        mask = _off_diagonal_mask(gram.size(-1), gram.device)
        off_diag = gram[:, mask].view(gram.size(0), -1)
        return off_diag.pow(2).mean()

    def _distribution_balance_loss(self, weights: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        flat_weights = weights.reshape(-1, weights.size(-1))
        flat_mask = mask.reshape(-1, mask.size(-1))
        valid = flat_mask.sum(dim=-1) > 1
        if not valid.any():
            return weights.new_zeros(())

        valid_weights = flat_weights[valid]
        valid_mask = flat_mask[valid]
        target = valid_mask / valid_mask.sum(dim=-1, keepdim=True).clamp_min(1.0)

        sample_loss = ((valid_weights - target).pow(2) * valid_mask).sum(dim=-1)
        sample_loss = sample_loss / valid_mask.sum(dim=-1).clamp_min(1.0)

        batch_mean = valid_weights.mean(dim=0)
        batch_target = target.mean(dim=0)
        batch_loss = (batch_mean - batch_target).pow(2).mean()
        return sample_loss.mean() + batch_loss

    def forward(self, outputs: Dict[str, torch.Tensor], labels: torch.Tensor) -> Dict[str, torch.Tensor]:
        logits = outputs["class_logits"]
        arcface_cosine = outputs["arcface_cosine"]
        prototype_sims = outputs["prototype_similarities"]
        branch_embeddings = outputs["branch_embeddings"]
        normalized_prototypes = outputs["normalized_prototypes"]
        embeddings = outputs["embedding"]
        branch_weights = outputs["branch_weights"]
        effective_branch_mask = outputs["effective_branch_mask"]
        stacked_view_weights = outputs["stacked_view_weights"]
        effective_view_mask = outputs["effective_view_mask"]

        classification = self._arcface_loss(arcface_cosine, labels)
        true_sims = prototype_sims[torch.arange(labels.size(0), device=labels.device), labels]
        compactness, assignments = self._compactness_loss(true_sims)
        diversity = self._prototype_diversity_loss(normalized_prototypes)
        usage = self._prototype_usage_loss(assignments, labels)
        branch_orthogonality = self._branch_orthogonality_loss(branch_embeddings)
        branch_balance = self._distribution_balance_loss(branch_weights, effective_branch_mask)
        expanded_view_mask = effective_view_mask.unsqueeze(1).expand_as(stacked_view_weights)
        view_balance = self._distribution_balance_loss(stacked_view_weights, expanded_view_mask)
        supcon = supervised_contrastive_loss(embeddings, labels, temperature=self.supcon_temperature)

        total = (
            self.classification_weight * classification
            + self.compactness_weight * compactness
            + self.diversity_weight * diversity
            + self.usage_weight * usage
            + self.branch_orthogonality_weight * branch_orthogonality
            + self.branch_balance_weight * branch_balance
            + self.view_balance_weight * view_balance
            + self.supcon_weight * supcon
        )

        return {
            "loss": total,
            "classification": classification.detach(),
            "compactness": compactness.detach(),
            "diversity": diversity.detach(),
            "usage": usage.detach(),
            "branch_orthogonality": branch_orthogonality.detach(),
            "branch_balance": branch_balance.detach(),
            "view_balance": view_balance.detach(),
            "supcon": supcon.detach(),
        }