forked from ccf-ai-infra/GPUCodeForces
53 lines
1.7 KiB
Python
53 lines
1.7 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
|
||
class Model(nn.Module):
|
||
"""
|
||
修正后的PyTorch ContrastiveLoss实现
|
||
使用标准公式：0.5 * [y * d² + (1-y) * max(0, margin - d)²]
|
||
其中 d = ||anchor - sample||₂ (实际欧氏距离)
|
||
"""
|
||
def __init__(self, margin=2.0):
|
||
super(Model, self).__init__()
|
||
self.margin = margin
|
||
|
||
def forward(self, anchor: torch.Tensor, sample: torch.Tensor, label: torch.Tensor) -> torch.Tensor:
|
||
"""
|
||
使用PyTorch标准实现ContrastiveLoss
|
||
|
||
Args:
|
||
anchor (torch.Tensor): 锚点样本 [batch_size, feature_dim]
|
||
sample (torch.Tensor): 对比样本 [batch_size, feature_dim]
|
||
label (torch.Tensor): 标签 [batch_size] (0=相似, 1=不相似)
|
||
|
||
Returns:
|
||
torch.Tensor: Contrastive Loss标量值
|
||
"""
|
||
# 使用PyTorch高度优化的pairwise_distance计算实际欧氏距离
|
||
distances = F.pairwise_distance(anchor, sample, p=2)
|
||
|
||
# 标准ContrastiveLoss公式
|
||
# y=1(不相似): 0.5 * max(0, margin - distance)²
|
||
# y=0(相似): 0.5 * distance²
|
||
losses = 0.5 * (label.float() * distances.pow(2) +
|
||
(1 - label).float() * F.relu(self.margin - distances).pow(2))
|
||
|
||
return torch.sum(losses)
|
||
|
||
batch_size = 128
|
||
feature_dim = 512
|
||
margin = 2.0
|
||
|
||
def get_inputs():
|
||
"""
|
||
生成合理的测试数据
|
||
"""
|
||
anchor = torch.randn(batch_size, feature_dim)
|
||
sample = torch.randn(batch_size, feature_dim)
|
||
label = torch.randint(0, 2, (batch_size,)).float()
|
||
return [anchor, sample, label]
|
||
|
||
def get_init_inputs():
|
||
return [margin] # margin参数
|