forked from ccf-ai-infra/GPUCodeForces
53 lines
1.6 KiB
Python
53 lines
1.6 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
|
||
class Model(nn.Module):
|
||
"""
|
||
Chebyshev + Sigmoid融合实现。
|
||
先计算Chebyshev距离，然后应用Sigmoid激活函数进行加权。
|
||
"""
|
||
def __init__(self):
|
||
super(Model, self).__init__()
|
||
|
||
def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||
"""
|
||
Compute Chebyshev + Sigmoid fusion.
|
||
|
||
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: Sigmoid-weighted Chebyshev distances [batch_size]
|
||
"""
|
||
# Input validation
|
||
if x.shape != y.shape:
|
||
raise ValueError(f"Input tensors must have the same shape, got {x.shape} and {y.shape}")
|
||
|
||
if x.dim() != 2:
|
||
raise ValueError(f"Input tensors must be 2D, got {x.dim()}D")
|
||
|
||
# Compute Chebyshev distance: max(|x_i - y_i|)
|
||
chebyshev_dist = torch.max(torch.abs(x - y), dim=1)[0]
|
||
|
||
# Apply Sigmoid activation
|
||
sigmoid_weights = torch.sigmoid(chebyshev_dist)
|
||
|
||
# Compute weighted distances
|
||
weighted_distances = sigmoid_weights * chebyshev_dist
|
||
|
||
return weighted_distances
|
||
|
||
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 [] # No special initialization inputs needed
|