forked from ccf-ai-infra/GPUCodeForces
33 lines
863 B
Python
33 lines
863 B
Python
import torch
|
||
import torch.nn as nn
|
||
|
||
|
||
class Model(nn.Module):
|
||
"""使用 PyTorch RMSNorm 的基准实现。"""
|
||
|
||
def __init__(self, weight: torch.Tensor, eps: float = 1e-6):
|
||
super().__init__()
|
||
if weight.dim() != 1:
|
||
raise ValueError("RMSNorm 权重必须是一维向量。")
|
||
feature_dim = weight.shape[0]
|
||
self.rmsnorm = nn.RMSNorm(feature_dim, eps=eps, elementwise_affine=True)
|
||
with torch.no_grad():
|
||
self.rmsnorm.weight.copy_(weight)
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
"""直接对输入做 RMSNorm，输出形状与输入一致。"""
|
||
return self.rmsnorm(x)
|
||
|
||
|
||
batch_size = 16
|
||
feature_dim = 2048
|
||
|
||
|
||
def get_inputs():
|
||
x = torch.randn(batch_size, feature_dim)
|
||
return [x]
|
||
|
||
|
||
def get_init_inputs():
|
||
weight = torch.randn(feature_dim)
|
||
return [weight] |