forked from ccf-ai-infra/GPUCodeForces
59 lines
1.9 KiB
Python
59 lines
1.9 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
|
||
class Model(nn.Module):
|
||
"""
|
||
PyTorch基准实现：Dice Loss + BCEWithLogitsLoss
|
||
"""
|
||
def __init__(self, dice_weight=1.0, bce_weight=1.0):
|
||
super(Model, self).__init__()
|
||
self.dice_weight = dice_weight
|
||
self.bce_weight = bce_weight
|
||
# BCEWithLogitsLoss内部融合了Sigmoid和BCE，更稳定
|
||
self.bce_loss_fn = nn.BCEWithLogitsLoss(reduction='mean')
|
||
|
||
def forward(self, logits: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
||
"""
|
||
Computes the combined Dice and BCE loss.
|
||
|
||
Args:
|
||
logits (torch.Tensor): Raw prediction tensor of shape [N, C, H, W].
|
||
target (torch.Tensor): Target tensor of same shape as logits.
|
||
|
||
Returns:
|
||
torch.Tensor: Combined loss value (scalar).
|
||
"""
|
||
# --- 第一步：计算Dice Loss ---
|
||
pred = torch.sigmoid(logits)
|
||
pred_flat = pred.view(-1)
|
||
target_flat = target.view(-1)
|
||
|
||
intersection = (pred_flat * target_flat).sum()
|
||
pred_sum = pred_flat.sum()
|
||
target_sum = target_flat.sum()
|
||
|
||
epsilon = 1e-6
|
||
dice_score = (2.0 * intersection) / (pred_sum + target_sum + epsilon)
|
||
loss_dice = 1.0 - dice_score
|
||
|
||
# --- 第二步：计算BCEWithLogitsLoss ---
|
||
# PyTorch的BCEWithLogitsLoss已经非常高效且稳定
|
||
loss_bce = self.bce_loss_fn(logits, target)
|
||
|
||
# --- 第三步：组合损失 ---
|
||
loss_total = self.dice_weight * loss_dice + self.bce_weight * loss_bce
|
||
|
||
return loss_total
|
||
|
||
batch_size = 32
|
||
height, width = 256, 256
|
||
channels = 1
|
||
|
||
def get_inputs():
|
||
logits = torch.randn(batch_size, channels, height, width)
|
||
target = torch.randint(0, 2, (batch_size, channels, height, width), dtype=torch.float32)
|
||
return [logits, target]
|
||
|
||
def get_init_inputs():
|
||
return [] # No special initialization inputs needed
|