forked from ccf-ai-infra/GPUCodeForces
49 lines
1.4 KiB
Python
49 lines
1.4 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
class Model(nn.Module):
|
|
"""
|
|
Simple model that performs Dice Loss calculation.
|
|
"""
|
|
def __init__(self):
|
|
super(Model, self).__init__()
|
|
|
|
def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Applies Dice Loss to the prediction and target tensors.
|
|
|
|
Args:
|
|
pred (torch.Tensor): Prediction tensor of any shape.
|
|
target (torch.Tensor): Target tensor of same shape as pred.
|
|
|
|
Returns:
|
|
torch.Tensor: Dice Loss value (scalar).
|
|
"""
|
|
# 展平张量
|
|
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()
|
|
|
|
# 计算Dice系数和损失
|
|
epsilon = 1e-6
|
|
dice_score = (2.0 * intersection) / (pred_sum + target_sum + epsilon)
|
|
dice_loss = 1.0 - dice_score
|
|
|
|
return dice_loss
|
|
|
|
batch_size = 32
|
|
height, width = 256, 256
|
|
channels = 1
|
|
|
|
def get_inputs():
|
|
pred = torch.rand(batch_size, channels, height, width)
|
|
target = torch.randint(0, 2, (batch_size, channels, height, width), dtype=torch.float32)
|
|
return [pred, target]
|
|
|
|
def get_init_inputs():
|
|
return [] # No special initialization inputs needed
|