forked from ccf-ai-infra/GPUCodeForces
75 lines
1.9 KiB
Python
75 lines
1.9 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
import math
|
|
|
|
# -------------------------------------------------------------
|
|
# 常量定义
|
|
# -------------------------------------------------------------
|
|
N_BATCH = 128
|
|
N_FEATURES = 512
|
|
|
|
# 损失函数参数
|
|
FULL = False
|
|
EPS = 1e-6
|
|
REDUCTION = 'mean'
|
|
|
|
|
|
# -------------------------------------------------------------
|
|
|
|
class Model(nn.Module):
|
|
"""
|
|
nn.GaussianNLLLoss 的纯 PyTorch 基准实现
|
|
"""
|
|
|
|
def __init__(self, full=False, eps=1e-6, reduction='mean'):
|
|
super().__init__()
|
|
self.full = full
|
|
self.eps = eps
|
|
self.reduction = reduction
|
|
|
|
if self.full:
|
|
self.const_term = 0.5 * math.log(2 * math.pi)
|
|
else:
|
|
self.const_term = 0.0
|
|
|
|
def forward(self, input: torch.Tensor, target: torch.Tensor, var: torch.Tensor) -> torch.Tensor:
|
|
|
|
# 1. 确保 var > eps。
|
|
# torch.clamp(min=...) 等价于 max(var, eps)
|
|
var_clamped = torch.clamp(var, min=self.eps)
|
|
|
|
# 2. 计算两个主要项
|
|
term1_log = torch.log(var_clamped)
|
|
term2_sq_err = (input - target).pow(2) / var_clamped
|
|
|
|
# 3. 组合
|
|
# (N, *) 形状
|
|
loss_unreduced = 0.5 * (term1_log + term2_sq_err) + self.const_term
|
|
|
|
# 4. 应用 Reduciton
|
|
if self.reduction == 'mean':
|
|
return loss_unreduced.mean()
|
|
elif self.reduction == 'sum':
|
|
return loss_unreduced.sum()
|
|
else: # 'none'
|
|
return loss_unreduced
|
|
|
|
|
|
def get_inputs():
|
|
"""
|
|
生成 (N, D) 形状的输入
|
|
"""
|
|
input = torch.randn(N_BATCH, N_FEATURES, dtype=torch.float32)
|
|
target = torch.randn(N_BATCH, N_FEATURES, dtype=torch.float32)
|
|
|
|
# Var 必须是正数
|
|
var = torch.rand(N_BATCH, N_FEATURES, dtype=torch.float32) + EPS
|
|
|
|
return [input, target, var]
|
|
|
|
|
|
def get_init_inputs():
|
|
return [FULL, EPS, REDUCTION]
|
|
|