forked from ccf-ai-infra/GPUCodeForces
41 lines
1.1 KiB
Python
41 lines
1.1 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
|
|
class Model(nn.Module):
|
|
def __init__(self, alpha=0.5, smooth=1e-6):
|
|
super().__init__()
|
|
self.alpha = alpha
|
|
self.smooth = smooth
|
|
|
|
def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
|
|
targets_f = targets.float()
|
|
|
|
probs = logits.sigmoid()
|
|
probs = probs.flatten(1)
|
|
targets_f = targets_f.flatten(1)
|
|
|
|
intersection = (probs * targets_f).sum(dim=1)
|
|
union = probs.sum(dim=1) + targets_f.sum(dim=1)
|
|
|
|
angle = 1.0 - (2.0 * intersection + self.smooth) / (union + self.smooth)
|
|
angular_loss = (angle * (1.0 - angle)).mean()
|
|
|
|
bce_loss = F.binary_cross_entropy_with_logits(logits, targets_f, reduction='mean')
|
|
|
|
return self.alpha * angular_loss + (1.0 - self.alpha) * bce_loss
|
|
|
|
|
|
batch_size = 512
|
|
feature_dim = 128
|
|
|
|
|
|
def get_inputs():
|
|
logits = torch.randn(batch_size, feature_dim, dtype=torch.float64)
|
|
targets = torch.randint(0, 2, (batch_size, feature_dim), dtype=torch.float64)
|
|
return [logits, targets]
|
|
|
|
|
|
def get_init_inputs():
|
|
return [0.5, 1e-6] |