forked from ccf-ai-infra/GPUCodeForces
60 lines
1.7 KiB
Python
60 lines
1.7 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
N, C, H, W = 32, 64, 56, 56
|
|
EPS = 1e-8
|
|
|
|
|
|
class CosineEmbeddingLossCustom(nn.Module):
|
|
|
|
def __init__(self, margin=0.0, reduction='mean', eps=1e-8):
|
|
super().__init__()
|
|
self.margin = margin
|
|
self.reduction = reduction
|
|
self.eps = eps
|
|
|
|
def forward(self, x1: torch.Tensor, x2: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
|
|
|
dot_product = torch.sum(x1 * x2, dim=1)
|
|
norm_x1 = torch.norm(x1, p=2, dim=1)
|
|
norm_x2 = torch.norm(x2, p=2, dim=1)
|
|
|
|
cos_sim = dot_product / (norm_x1 * norm_x2 + self.eps)
|
|
|
|
loss_pos = 1.0 - cos_sim
|
|
loss_neg = F.relu(cos_sim - self.margin)
|
|
|
|
loss = torch.where(target == 1, loss_pos, loss_neg)
|
|
|
|
if self.reduction == 'mean':
|
|
return loss.mean()
|
|
elif self.reduction == 'sum':
|
|
return loss.sum()
|
|
else:
|
|
return loss
|
|
|
|
|
|
class Model(nn.Module):
|
|
def __init__(self, margin=0.5):
|
|
super().__init__()
|
|
self.op = CosineEmbeddingLossCustom(margin=margin, reduction='mean', eps=EPS)
|
|
|
|
def forward(self, x1: torch.Tensor, x2: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
|
x1_flat = x1.view(x1.size(0), -1)
|
|
x2_flat = x2.view(x2.size(0), -1)
|
|
return self.op(x1_flat, x2_flat, target)
|
|
|
|
|
|
def get_inputs():
|
|
x1 = torch.randn(N, C, H, W, dtype=torch.float32)
|
|
x2 = torch.randn(N, C, H, W, dtype=torch.float32)
|
|
|
|
target = torch.randint(0, 2, (N,), dtype=torch.float32) # 0 or 1
|
|
target = torch.where(target == 0, torch.tensor(-1.0), torch.tensor(1.0))
|
|
|
|
return [x1, x2, target]
|
|
|
|
|
|
def get_init_inputs():
|
|
return [0.5] |