forked from ccf-ai-infra/GPUCodeForces
50 lines
1.3 KiB
Python
50 lines
1.3 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
N, C, H, W = 32, 64, 56, 56
|
|
|
|
|
|
class SmoothL1Loss(nn.Module):
|
|
def __init__(self, reduction='mean', beta=1.0):
|
|
super().__init__()
|
|
self.reduction = reduction
|
|
self.beta = beta
|
|
|
|
def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
|
diff = torch.abs(input - target)
|
|
|
|
if self.beta == 0:
|
|
loss = diff
|
|
else:
|
|
loss = torch.where(
|
|
diff < self.beta,
|
|
0.5 * diff * diff / self.beta,
|
|
diff - 0.5 * self.beta
|
|
)
|
|
|
|
if self.reduction == 'mean':
|
|
return loss.mean()
|
|
elif self.reduction == 'sum':
|
|
return loss.sum()
|
|
else:
|
|
return loss
|
|
|
|
|
|
class Model(nn.Module):
|
|
def __init__(self, reduction='mean', beta=1.0):
|
|
super().__init__()
|
|
self.op = SmoothL1Loss(reduction, beta)
|
|
|
|
def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
|
return self.op(input, target)
|
|
|
|
|
|
def get_inputs():
|
|
input = torch.randn(N, C, H, W, dtype=torch.float32)
|
|
target = torch.randn(N, C, H, W, dtype=torch.float32)
|
|
return [input, target]
|
|
|
|
|
|
def get_init_inputs():
|
|
return ['mean', 1.0] |