28 lines
750 B
Python
28 lines
750 B
Python
import torch.nn as nn
|
|
|
|
|
|
class LinearHead(nn.Module):
|
|
def __init__(self, input_dim, output_dim):
|
|
super().__init__()
|
|
self.norm = nn.LayerNorm(input_dim)
|
|
self.linear = nn.Linear(input_dim, output_dim)
|
|
nn.init.normal_(self.linear.weight, std=0.02)
|
|
nn.init.zeros_(self.linear.bias)
|
|
|
|
def forward(self, x):
|
|
return self.linear(self.norm(x))
|
|
|
|
|
|
class MLPBlock(nn.Module):
|
|
def __init__(self, input_dim=512, hidden_dim=256, dropout=0.1):
|
|
super().__init__()
|
|
self.net = nn.Sequential(
|
|
nn.LayerNorm(input_dim),
|
|
nn.Linear(input_dim, hidden_dim),
|
|
nn.GELU(),
|
|
nn.Dropout(dropout),
|
|
)
|
|
|
|
def forward(self, x):
|
|
return self.net(x)
|