forked from ccf-ai-infra/GPUCodeForces
55 lines
1.6 KiB
Python
55 lines
1.6 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class Model(nn.Module):
|
|
def __init__(self, bins=32, eps=1e-12):
|
|
super().__init__()
|
|
self.bins = bins
|
|
self.eps = eps
|
|
|
|
def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
|
b, n = x.shape
|
|
|
|
min_x = x.min(dim=1, keepdim=True)[0]
|
|
max_x = x.max(dim=1, keepdim=True)[0]
|
|
min_y = y.min(dim=1, keepdim=True)[0]
|
|
max_y = y.max(dim=1, keepdim=True)[0]
|
|
|
|
x_norm = (x - min_x) / (max_x - min_x + 1e-6)
|
|
y_norm = (y - min_y) / (max_y - min_y + 1e-6)
|
|
|
|
x_bin = torch.clamp((x_norm * self.bins).long(), 0, self.bins - 1)
|
|
y_bin = torch.clamp((y_norm * self.bins).long(), 0, self.bins - 1)
|
|
|
|
joint_idx = x_bin * self.bins + y_bin
|
|
hist = torch.zeros(b, self.bins * self.bins, device=x.device, dtype=torch.float32)
|
|
ones = torch.ones_like(joint_idx, dtype=torch.float32)
|
|
hist.scatter_add_(1, joint_idx, ones)
|
|
|
|
p_xy = hist.view(b, self.bins, self.bins) / n
|
|
|
|
p_x = p_xy.sum(dim=2) # (B, bins)
|
|
p_y = p_xy.sum(dim=1) # (B, bins)
|
|
|
|
h_x = -torch.sum(p_x * torch.log(p_x + self.eps), dim=1)
|
|
h_y = -torch.sum(p_y * torch.log(p_y + self.eps), dim=1)
|
|
|
|
p_x_p_y = torch.bmm(p_x.unsqueeze(2), p_y.unsqueeze(1))
|
|
mi = torch.sum(p_xy * torch.log((p_xy + self.eps) / (p_x_p_y + self.eps)), dim=(1, 2))
|
|
|
|
return h_x + h_y - 2 * mi
|
|
|
|
|
|
batch_size = 128
|
|
feature_dim = 512
|
|
|
|
|
|
def get_inputs():
|
|
x = torch.randn(batch_size, feature_dim, dtype=torch.float32)
|
|
y = torch.randn(batch_size, feature_dim, dtype=torch.float32)
|
|
return [x, y]
|
|
|
|
|
|
def get_init_inputs():
|
|
return [32] |