forked from ccf-ai-infra/GPUCodeForces
55 lines
1.7 KiB
Python
55 lines
1.7 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
N, D = 32, 128
|
|
|
|
|
|
class TripletMarginWithDistanceLoss(nn.Module):
|
|
|
|
def __init__(self, distance_function=None, margin=1.0, swap=False, reduction='mean'):
|
|
super().__init__()
|
|
self.distance_function = distance_function if distance_function is not None else nn.PairwiseDistance()
|
|
self.margin = margin
|
|
self.swap = swap
|
|
self.reduction = reduction
|
|
|
|
def forward(self, anchor: torch.Tensor, positive: torch.Tensor, negative: torch.Tensor) -> torch.Tensor:
|
|
|
|
d_ap = self.distance_function(anchor, positive)
|
|
|
|
d_an = self.distance_function(anchor, negative)
|
|
|
|
if self.swap:
|
|
d_pn = self.distance_function(positive, negative)
|
|
d_an = torch.min(d_an, d_pn)
|
|
|
|
loss = torch.clamp(d_ap - d_an + self.margin, min=0.0)
|
|
|
|
if self.reduction == 'mean':
|
|
return loss.mean()
|
|
elif self.reduction == 'sum':
|
|
return loss.sum()
|
|
else: # 'none'
|
|
return loss
|
|
|
|
|
|
class Model(nn.Module):
|
|
def __init__(self, margin=1.0, swap=False):
|
|
super().__init__()
|
|
|
|
self.op = TripletMarginWithDistanceLoss(distance_function=nn.PairwiseDistance(), margin=margin, swap=swap)
|
|
|
|
def forward(self, a: torch.Tensor, p: torch.Tensor, n: torch.Tensor) -> torch.Tensor:
|
|
return self.op(a, p, n)
|
|
|
|
|
|
def get_inputs():
|
|
anchor = torch.randn(N, D, dtype=torch.float32)
|
|
positive = torch.randn(N, D, dtype=torch.float32)
|
|
negative = torch.randn(N, D, dtype=torch.float32)
|
|
return [anchor, positive, negative]
|
|
|
|
|
|
def get_init_inputs():
|
|
return [1.0, False] |