forked from ccf-ai-infra/GPUCodeForces
39 lines
1.2 KiB
Python
39 lines
1.2 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
class Model(nn.Module):
|
|
"""
|
|
Baseline: Cosine Similarity + Contrastive Loss
|
|
"""
|
|
def __init__(self, margin=0.5):
|
|
super(Model, self).__init__()
|
|
self.margin = margin
|
|
|
|
def forward(self, x: torch.Tensor, y: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
|
|
# 1. 计算余弦相似度
|
|
norm_x = torch.sqrt(torch.sum(x * x, dim=1, keepdim=True))
|
|
norm_y = torch.sqrt(torch.sum(y * y, dim=1, keepdim=True))
|
|
dot_product = torch.sum(x * y, dim=1, keepdim=True)
|
|
cosine_sim = dot_product / (norm_x * norm_y + 1e-8)
|
|
cosine_sim = cosine_sim.squeeze(1)
|
|
|
|
# 2. 计算对比损失
|
|
loss_positive = labels * (1 - cosine_sim)
|
|
loss_negative = (1 - labels) * torch.relu(cosine_sim - self.margin)
|
|
loss = loss_positive + loss_negative
|
|
|
|
return torch.sum(loss)
|
|
|
|
batch_size = 1024
|
|
feature_dim = 512
|
|
|
|
def get_inputs():
|
|
x = torch.randn(batch_size, feature_dim)
|
|
y = torch.randn(batch_size, feature_dim)
|
|
# 生成0或1的标签
|
|
labels = torch.randint(0, 2, (batch_size,)).float()
|
|
return [x, y, labels]
|
|
|
|
def get_init_inputs():
|
|
return []
|