forked from ccf-ai-infra/GPUCodeForces
40 lines
958 B
Python
40 lines
958 B
Python
import torch
|
|
import torch.nn as nn
|
|
import math
|
|
|
|
N, C, H, W = 32, 64, 56, 56
|
|
|
|
|
|
class LogCoshLoss(nn.Module):
|
|
def __init__(self, reduction='mean'):
|
|
super().__init__()
|
|
self.reduction = reduction
|
|
|
|
def forward(self, input, target):
|
|
diff = input - target
|
|
loss = torch.abs(diff) + torch.nn.functional.softplus(-2. * torch.abs(diff)) - math.log(2.0)
|
|
|
|
if self.reduction == 'mean':
|
|
return loss.mean()
|
|
elif self.reduction == 'sum':
|
|
return loss.sum()
|
|
return loss
|
|
|
|
|
|
class Model(nn.Module):
|
|
def __init__(self, reduction='mean'):
|
|
super().__init__()
|
|
self.op = LogCoshLoss(reduction)
|
|
|
|
def forward(self, input, target):
|
|
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'] |