forked from ccf-ai-infra/GPUCodeForces
28 lines
681 B
Python
28 lines
681 B
Python
# cosineloss_torch.py
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
|
|
BATCH_SIZE = 16
|
|
EMBEDDING_DIM = 256
|
|
DIM = BATCH_SIZE
|
|
|
|
MARGIN = 0.5
|
|
|
|
class Model(nn.Module):
|
|
def forward(self, x1: torch.Tensor, x2: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
|
return F.cosine_embedding_loss(x1, x2, y, margin=MARGIN, reduction='mean')
|
|
|
|
def get_inputs():
|
|
|
|
x1 = torch.randn(BATCH_SIZE, EMBEDDING_DIM, dtype=torch.float32)
|
|
x2 = torch.randn(BATCH_SIZE, EMBEDDING_DIM, dtype=torch.float32)
|
|
|
|
y = torch.randint(0, 2, size=(BATCH_SIZE,), dtype=torch.float32)
|
|
y[y == 0] = -1.0
|
|
|
|
return [x1, x2, y]
|
|
|
|
def get_init_inputs():
|
|
return [] |