| 2 | import torch.nn as nn |
| 3 | |
| 4 | class LayerNormalization(nn.Module): |
| 5 | def __init__(self, dim_embed, epsilon=1e-6): |
| 6 | super(LayerNormalization, self).__init__() |
| 7 | self.epsilon = epsilon |
| 8 | self.gamma = nn.Parameter(torch.ones(dim_embed)) |
| 9 | self.beta = nn.Parameter(torch.zeros(dim_embed)) |
| 10 | |
| 11 | def forward(self, x): |
| 12 | mean = x.mean(dim=-1, keepdim=True) |
| 13 | std = x.std(dim=-1, keepdim=True) |
| 14 | normalized_x = (x - mean) / (std + self.epsilon) |
| 15 | return self.gamma * normalized_x + self.beta |
| 16 | |
| 17 | |
| 18 | if __name__ == "__main__": |