forked from ccf-ai-infra/GPUCodeForces
52 lines
1.6 KiB
Python
52 lines
1.6 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
|
||
class Model(nn.Module):
|
||
"""
|
||
Simple model that performs LayerNorm normalization using PyTorch's built-in nn.LayerNorm.
|
||
"""
|
||
|
||
def __init__(self, normalized_shape=None, eps=1e-5, elementwise_affine=True):
|
||
super(Model, self).__init__()
|
||
# 如果未指定normalized_shape，将在forward中动态设置
|
||
self.normalized_shape = normalized_shape
|
||
self.eps = eps
|
||
self.elementwise_affine = elementwise_affine
|
||
self.layernorm = None
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
"""
|
||
Applies LayerNorm to the input tensor.
|
||
|
||
Args:
|
||
x (torch.Tensor): Input tensor of any shape.
|
||
|
||
Returns:
|
||
torch.Tensor: Output tensor with LayerNorm applied, same shape as input.
|
||
"""
|
||
# 如果layernorm未初始化，根据输入形状动态创建
|
||
if self.layernorm is None:
|
||
if self.normalized_shape is None:
|
||
# 默认对最后一个维度进行归一化
|
||
self.normalized_shape = x.shape[1:]
|
||
self.layernorm = nn.LayerNorm(
|
||
normalized_shape=self.normalized_shape,
|
||
eps=self.eps,
|
||
elementwise_affine=self.elementwise_affine
|
||
).to(x.device)
|
||
|
||
return self.layernorm(x)
|
||
|
||
|
||
batch_size = 16
|
||
dim = 16384
|
||
|
||
|
||
def get_inputs():
|
||
x = torch.randn(batch_size, dim)
|
||
return [x]
|
||
|
||
|
||
def get_init_inputs():
|
||
# 可以传入归一化形状、eps等参数，保持向后兼容
|
||
return [] # 使用默认参数 |