forked from ccf-ai-infra/GPUCodeForces
52 lines
1.5 KiB
Python
52 lines
1.5 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
|
||
class Model(nn.Module):
|
||
"""
|
||
Huber Loss implementation - robust loss function that is quadratic for small errors and linear for large errors.
|
||
"""
|
||
def __init__(self, delta=1.0):
|
||
super(Model, self).__init__()
|
||
self.delta = delta
|
||
|
||
def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
||
"""
|
||
Compute Huber loss between input and target tensors.
|
||
|
||
Args:
|
||
input (torch.Tensor): Predicted values [batch_size, ...]
|
||
target (torch.Tensor): Target values [batch_size, ...]
|
||
|
||
Returns:
|
||
torch.Tensor: Scalar Huber loss value
|
||
"""
|
||
# 计算绝对误差
|
||
abs_error = torch.abs(input - target)
|
||
|
||
# 创建二次区域和线性区域的mask
|
||
quadratic_mask = abs_error <= self.delta
|
||
linear_mask = ~quadratic_mask
|
||
|
||
# 二次区域：0.5 * error²
|
||
quadratic_loss = 0.5 * torch.pow(input - target, 2)
|
||
|
||
# 线性区域：delta * (|error| - 0.5 * delta)
|
||
linear_loss = self.delta * (abs_error - 0.5 * self.delta)
|
||
|
||
# 组合损失
|
||
loss = torch.where(quadratic_mask, quadratic_loss, linear_loss)
|
||
|
||
return torch.sum(loss)
|
||
|
||
batch_size = 512
|
||
feature_dim = 512
|
||
delta = 1.0
|
||
|
||
def get_inputs():
|
||
input = torch.randn(batch_size, feature_dim)
|
||
target = torch.randn(batch_size, feature_dim)
|
||
return [input, target]
|
||
|
||
def get_init_inputs():
|
||
return [delta] # delta parameter
|