mirror of
https://github.com/wassname/Run-Skeleton-Run.git
synced 2026-06-27 19:45:48 +08:00
16 lines
446 B
Python
16 lines
446 B
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class LayerNorm(nn.Module):
|
|
def __init__(self, features, eps=1e-6):
|
|
super().__init__()
|
|
self.gamma = nn.Parameter(torch.ones(features))
|
|
self.beta = nn.Parameter(torch.zeros(features))
|
|
self.eps = eps
|
|
|
|
def forward(self, x):
|
|
mean = x.mean(-1, keepdim=True)
|
|
std = x.std(-1, keepdim=True)
|
|
return self.gamma * (x - mean) / (std + self.eps) + self.beta
|