forked from ccf-ai-infra/GPUCodeForces
45 lines
1.3 KiB
Python
45 lines
1.3 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
|
||
BATCH_SIZE = 4096
|
||
FEATURE_DIM = 512
|
||
|
||
# --- 损失函数的参数 ---
|
||
# log_input=True: loss = exp(input) - target * input
|
||
# log_input=False: loss = input - target * log(input + eps)
|
||
LOG_INPUT = True
|
||
# full=True: 添加 Stirling's approximation
|
||
FULL = False
|
||
EPS = 1e-8
|
||
|
||
|
||
class Model(nn.Module):
|
||
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.criterion = nn.PoissonNLLLoss(
|
||
log_input=LOG_INPUT,
|
||
full=FULL,
|
||
eps=EPS,
|
||
reduction='mean'
|
||
)
|
||
|
||
def forward(self, input_tensor: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
||
# 在 PyTorch 中，input_tensor 是文档中的 'input'
|
||
return self.criterion(input_tensor, target)
|
||
|
||
|
||
def get_inputs():
|
||
# Input (log_input=True 时) 可以是任意实数
|
||
input_tensor = torch.randn(BATCH_SIZE, FEATURE_DIM, dtype=torch.float32)
|
||
|
||
# Target 在 Poisson 分布中代表计数，且在 'full' 模式下会计算 log(target)
|
||
# 因此 target 必须是 >= 0 的。我们使用 rand 来确保
|
||
target = torch.rand(BATCH_SIZE, FEATURE_DIM, dtype=torch.float32) * 10 # 乘以 10 以便有一些 > 1
|
||
|
||
return [input_tensor, target]
|
||
|
||
|
||
def get_init_inputs():
|
||
return [] |