forked from ccf-ai-infra/GPUCodeForces
46 lines
1.3 KiB
Python
46 lines
1.3 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
class Model(nn.Module):
|
|
"""
|
|
Cosine Similarity implementation.
|
|
Computes the cosine similarity between two sets of vectors.
|
|
"""
|
|
def __init__(self):
|
|
super(Model, self).__init__()
|
|
|
|
def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Compute cosine similarity between x and y.
|
|
|
|
Args:
|
|
x (torch.Tensor): First set of vectors [batch_size, feature_dim]
|
|
y (torch.Tensor): Second set of vectors [batch_size, feature_dim]
|
|
|
|
Returns:
|
|
torch.Tensor: Cosine similarities [batch_size]
|
|
"""
|
|
# 计算L2范数
|
|
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)
|
|
|
|
return cosine_sim.squeeze(1)
|
|
|
|
batch_size = 1024
|
|
feature_dim = 512
|
|
|
|
def get_inputs():
|
|
# Generate two sets of vectors
|
|
x = torch.randn(batch_size, feature_dim)
|
|
y = torch.randn(batch_size, feature_dim)
|
|
return [x, y]
|
|
|
|
def get_init_inputs():
|
|
return [] # No special initialization inputs needed
|