forked from ccf-ai-infra/GPUCodeForces
47 lines
1.4 KiB
Python
47 lines
1.4 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
N, C, H, W = 32, 1, 64, 64
|
|
|
|
|
|
class GeneratorLoss(nn.Module):
|
|
def __init__(self, reduction='mean', beta=1.0):
|
|
super().__init__()
|
|
self.reduction = reduction
|
|
if reduction not in ['none', 'mean', 'sum']:
|
|
raise ValueError("Invalid reduction mode")
|
|
|
|
def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
|
|
|
loss = F.relu(input) - input * target + torch.log1p(torch.exp(-torch.abs(input)))
|
|
|
|
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 = GeneratorLoss(reduction, beta)
|
|
|
|
def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
|
if isinstance(input, (list, tuple)) and len(input) > 0:
|
|
input = input[0]
|
|
target = target[0] if len(target) > 0 else target
|
|
|
|
return self.op(input, target)
|
|
|
|
|
|
def get_inputs():
|
|
input = torch.randn(N, C, H, W, dtype=torch.float32)
|
|
target = torch.ones(N, C, H, W, dtype=torch.float32)
|
|
return [input, target]
|
|
|
|
|
|
def get_init_inputs():
|
|
return ['mean', 1.0] |