forked from ccf-ai-infra/GPUCodeForces
95 lines
2.5 KiB
Python
95 lines
2.5 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
# -------------------------------------------------------------
|
|
# 常量定义
|
|
# -------------------------------------------------------------
|
|
N, C, H, W = 8, 10, 16, 16 # (N, C, H, W)
|
|
|
|
# 损失函数参数
|
|
WEIGHT = torch.rand(C, dtype=torch.float32) # (C,)
|
|
IGNORE_INDEX = -100
|
|
REDUCTION = 'mean'
|
|
|
|
|
|
# -------------------------------------------------------------
|
|
|
|
class Model(nn.Module):
|
|
"""
|
|
nn.NLLLoss 的纯 PyTorch 基准实现
|
|
(K-dim, 2D-example)
|
|
"""
|
|
|
|
def __init__(self, weight=None, size_average=None, ignore_index=-100,
|
|
reduce=None, reduction='mean'):
|
|
super().__init__()
|
|
|
|
# 处理已弃用的 size_average 和 reduce
|
|
if size_average is not None or reduce is not None:
|
|
# (省略... 遵循 torch.nn.modules.loss)
|
|
pass
|
|
|
|
self.reduction = reduction
|
|
self.ignore_index = ignore_index
|
|
|
|
# 确保 weight 在正确的设备上
|
|
if weight is not None:
|
|
self.register_buffer('weight', weight)
|
|
else:
|
|
self.weight = None
|
|
|
|
def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
|
|
|
input_flat = input.view(N, C, -1)
|
|
|
|
target_flat = target.view(N, -1)
|
|
|
|
|
|
loss_unreduced = input_flat.gather(dim=1, index=target_flat.unsqueeze(1))
|
|
loss_unreduced = -loss_unreduced.squeeze(1) # (N, H*W)
|
|
|
|
|
|
if self.weight is not None:
|
|
|
|
weights_applied = self.weight[target_flat]
|
|
loss_unreduced = loss_unreduced * weights_applied
|
|
else:
|
|
|
|
weights_applied = torch.ones_like(target_flat, dtype=input.dtype)
|
|
|
|
|
|
mask = (target_flat != self.ignore_index)
|
|
loss_unreduced = loss_unreduced * mask
|
|
weights_applied = weights_applied * mask
|
|
|
|
|
|
if self.reduction == 'mean':
|
|
|
|
total_weight = weights_applied.sum()
|
|
if total_weight == 0:
|
|
return torch.tensor(0.0, device=input.device, dtype=input.dtype)
|
|
return loss_unreduced.sum() / total_weight
|
|
|
|
elif self.reduction == 'sum':
|
|
return loss_unreduced.sum()
|
|
|
|
else:
|
|
return loss_unreduced.view_as(target)
|
|
|
|
|
|
def get_inputs():
|
|
|
|
input_log_probs = F.log_softmax(torch.randn(N, C, H, W, dtype=torch.float32), dim=1)
|
|
target = torch.empty(N, H, W, dtype=torch.long).random_(0, C)
|
|
|
|
|
|
target.view(-1)[::10] = IGNORE_INDEX
|
|
|
|
return [input_log_probs, target]
|
|
|
|
|
|
def get_init_inputs():
|
|
|
|
return [WEIGHT, None, IGNORE_INDEX, None, REDUCTION]
|