forked from ccf-ai-infra/GPUCodeForces
60 lines
1.9 KiB
Python
60 lines
1.9 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
# 定义常量
|
|
BATCH_SIZE = 256
|
|
FEATURE_DIM = 512
|
|
MARGIN = 0.25
|
|
GAMMA = 256
|
|
|
|
|
|
class Model(nn.Module):
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.margin = MARGIN
|
|
self.gamma = GAMMA
|
|
|
|
def forward(self, features: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
|
|
# Circle Loss的 PyTorch 实现
|
|
# 假设 features 已经是 L2 归一化的
|
|
# (B, D) @ (D, B) -> (B, B)
|
|
similarities = torch.matmul(features, features.t())
|
|
|
|
# 创建正负样本对的掩码
|
|
mask_positive = labels.unsqueeze(1) == labels.unsqueeze(0)
|
|
mask_negative = labels.unsqueeze(1) != labels.unsqueeze(0)
|
|
|
|
# 收集正样本对和负样本对的相似度
|
|
# .masked_select() 会将张量展平
|
|
sp = similarities[mask_positive]
|
|
sn = similarities[mask_negative]
|
|
|
|
# 计算 Circle Loss 的 logits
|
|
# .detach() 用于停止梯度反向传播
|
|
ap = torch.clamp_min(-sp.detach() + 1 + self.margin, min=0.)
|
|
an = torch.clamp_min(sn.detach() + self.margin, min=0.)
|
|
|
|
delta_p = 1 - self.margin
|
|
delta_n = self.margin
|
|
|
|
logit_p = -ap * (sp - delta_p) * self.gamma
|
|
logit_n = an * (sn - delta_n) * self.gamma
|
|
|
|
# 使用 logsumexp 和 softplus 计算最终的 loss
|
|
# 这是 "unified" 版本的 loss
|
|
loss = F.softplus(torch.logsumexp(logit_n, dim=0) + torch.logsumexp(logit_p, dim=0))
|
|
|
|
return loss
|
|
|
|
|
|
def get_inputs():
|
|
# 特征需要 L2 归一化
|
|
features = F.normalize(torch.randn(BATCH_SIZE, FEATURE_DIM, dtype=torch.float32), p=2, dim=1)
|
|
labels = torch.randint(0, 10, (BATCH_SIZE,), dtype=torch.long) # 假设有 10 个类别
|
|
return [features, labels]
|
|
|
|
|
|
def get_init_inputs():
|
|
return [] |