forked from ccf-ai-infra/GPUCodeForces
39 lines
1.1 KiB
Python
39 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=1.0):
|
|
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_flat = targets_f.flatten(1)
|
|
|
|
intersection = (probs * targets_f_flat).sum(dim=1)
|
|
dice = (2.0 * intersection + self.smooth) / (probs.sum(dim=1) + targets_f_flat.sum(dim=1) + self.smooth)
|
|
dice_loss = 1.0 - dice.mean()
|
|
|
|
ce_loss = F.binary_cross_entropy_with_logits(logits, targets_f, reduction='mean')
|
|
|
|
return self.alpha * dice_loss + (1.0 - self.alpha) * ce_loss
|
|
|
|
|
|
batch_size = 128
|
|
feature_dim = 100
|
|
|
|
|
|
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, 1.0] |