forked from ccf-ai-infra/GPUCodeForces
59 lines
1.6 KiB
Python
59 lines
1.6 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
|
|
class Model(nn.Module):
|
|
def __init__(self, num_bins=10):
|
|
super(Model, self).__init__()
|
|
self.num_bins = num_bins
|
|
self.max_dist = 4.0
|
|
|
|
def forward(self, x: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
|
|
x = F.normalize(x, p=2, dim=1)
|
|
|
|
dist = torch.cdist(x, x, p=2).pow(2)
|
|
|
|
width = self.max_dist / self.num_bins
|
|
centers = torch.linspace(width / 2, self.max_dist - width / 2, self.num_bins, device=x.device)
|
|
|
|
d = dist.unsqueeze(-1)
|
|
c = centers.view(1, 1, -1)
|
|
|
|
bin_weights = torch.relu(1.0 - torch.abs(d - c) / width)
|
|
|
|
eq_labels = labels.unsqueeze(1) == labels.unsqueeze(0)
|
|
diag_mask = torch.eye(x.size(0), device=x.device, dtype=torch.bool)
|
|
|
|
pos_mask = eq_labels & (~diag_mask)
|
|
neg_mask = (~eq_labels) & (~diag_mask)
|
|
|
|
pos_hist = (bin_weights * pos_mask.unsqueeze(-1)).sum(dim=1)
|
|
neg_hist = (bin_weights * neg_mask.unsqueeze(-1)).sum(dim=1)
|
|
|
|
pos_cdf = torch.cumsum(pos_hist, dim=1)
|
|
neg_cdf = torch.cumsum(neg_hist, dim=1)
|
|
|
|
precision = pos_cdf / (pos_cdf + neg_cdf + 1e-10)
|
|
|
|
total_pos = pos_cdf[:, -1].unsqueeze(1) + 1e-10
|
|
delta_recall = pos_hist / total_pos
|
|
|
|
ap = (precision * delta_recall).sum(dim=1)
|
|
|
|
return 1.0 - ap.mean()
|
|
|
|
|
|
batch_size = 128
|
|
input_dim = 1024
|
|
num_classes = 32
|
|
|
|
|
|
def get_inputs():
|
|
x = torch.randn(batch_size, input_dim)
|
|
labels = torch.randint(0, num_classes, (batch_size,))
|
|
return [x, labels]
|
|
|
|
|
|
def get_init_inputs():
|
|
return [] |