22 lines
623 B
Python
22 lines
623 B
Python
import torch.nn as nn
|
|
|
|
def init_linear(l, stddev):
|
|
nn.init.normal_(l.weight, std=stddev)
|
|
if l.bias is not None:
|
|
nn.init.constant_(l.bias, 0.0)
|
|
|
|
class MLP(nn.Module):
|
|
def __init__(self, *,
|
|
width: int,
|
|
init_scale: float):
|
|
super().__init__()
|
|
self.width = width
|
|
self.c_fc = nn.Linear(width, width * 4)
|
|
self.c_proj = nn.Linear(width * 4, width)
|
|
self.gelu = nn.GELU()
|
|
init_linear(self.c_fc, init_scale)
|
|
init_linear(self.c_proj, init_scale)
|
|
|
|
def forward(self, x):
|
|
return self.c_proj(self.gelu(self.c_fc(x)))
|