forked from ccf-ai-infra/GPUCodeForces
32 lines
1002 B
Python
32 lines
1002 B
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
import math
|
|
|
|
|
|
class ModelNew(nn.Module):
|
|
def __init__(self, in_features: int = 1024, out_features: int = 2048):
|
|
super().__init__()
|
|
self.in_features = in_features
|
|
self.out_features = out_features
|
|
self.weight = nn.Parameter(torch.empty(out_features, in_features))
|
|
self.bias = nn.Parameter(torch.zeros(out_features))
|
|
nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
|
|
fan_in = self.weight.size(1)
|
|
bound = 1.0 / math.sqrt(fan_in)
|
|
nn.init.uniform_(self.bias, -bound, bound)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
y = F.linear(x, self.weight, self.bias)
|
|
# 使用精确 GELU 以确保与基线一致的数值结果
|
|
return F.gelu(y, approximate='none')
|
|
|
|
|
|
def get_init_inputs():
|
|
return {"in_features": 1024, "out_features": 2048}
|
|
|
|
|
|
def get_inputs():
|
|
B, T, D = 16, 512, 1024
|
|
x = torch.randn(B, T, D)
|
|
return x |