forked from ccf-ai-infra/GPUCodeForces
34 lines
924 B
Python
34 lines
924 B
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
class Model(nn.Module):
|
|
def __init__(self, hidden_size):
|
|
super().__init__()
|
|
torch.manual_seed(42)
|
|
self.hidden_size = hidden_size
|
|
self.weight = nn.Parameter(torch.ones(hidden_size))
|
|
self.bias = nn.Parameter(torch.zeros(hidden_size))
|
|
self.eps = 1e-5
|
|
|
|
def forward(self, x):
|
|
swish_out = x * torch.sigmoid(x)
|
|
|
|
mean = swish_out.mean(-1, keepdim=True)
|
|
var = swish_out.var(-1, keepdim=True, unbiased=False)
|
|
swish_out = (swish_out - mean) / torch.sqrt(var + self.eps)
|
|
|
|
output = swish_out * self.weight + self.bias
|
|
|
|
return output
|
|
|
|
def get_inputs():
|
|
batch_size = 4096
|
|
seq_len = 128
|
|
hidden_size = 768 # 典型的BERT hidden size
|
|
x = torch.randn(batch_size, seq_len, hidden_size)
|
|
return [x]
|
|
|
|
def get_init_inputs():
|
|
return [768]
|