forked from ccf-ai-infra/GPUCodeForces
45 lines
1.2 KiB
Python
45 lines
1.2 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
class Model(nn.Module):
|
|
"""
|
|
Cosine Distance implementation - 标准版本
|
|
"""
|
|
def __init__(self, eps=1e-8):
|
|
super(Model, self).__init__()
|
|
self.eps = eps
|
|
|
|
def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Compute cosine distance between x and y.
|
|
Cosine Distance = 1 - Cosine Similarity
|
|
|
|
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 distances [batch_size], range [0, 2]
|
|
"""
|
|
# 使用PyTorch内置的cosine_similarity函数
|
|
cosine_sim = F.cosine_similarity(x, y, dim=1, eps=self.eps)
|
|
|
|
# 转换为距离
|
|
cosine_dist = 1.0 - cosine_sim
|
|
|
|
# 确保结果在合理范围内
|
|
return torch.clamp(cosine_dist, 0.0, 2.0)
|
|
|
|
batch_size = 256
|
|
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 [1e-8] # eps value
|