Spaces:
Runtime error
Runtime error
File size: 1,564 Bytes
c3d0293 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 |
from torch import nn
import torch
import torch.nn.functional as F
class LayerNorm(nn.Module):
def __init__(self, normalized_shape, eps = 1e-5, elementwise_affine = True,
device=None, dtype=None):
factory_kwargs = {'device': device, 'dtype': dtype}
super().__init__()
if isinstance(normalized_shape, int):
normalized_shape = [normalized_shape]
self.normalized_shape = normalized_shape # type: ignore[arg-type]
self.eps = eps
self.elementwise_affine = elementwise_affine
if self.elementwise_affine:
self.weight = nn.parameter.Parameter(torch.ones(self.normalized_shape, **factory_kwargs))
self.bias = nn.parameter.Parameter(torch.zeros(self.normalized_shape, **factory_kwargs))
else:
self.register_parameter('weight', None)
self.register_parameter('bias', None)
def forward(self, input):
orig_type = input.dtype
ret = F.layer_norm(input.type(torch.float32), self.normalized_shape, self.weight.type(torch.float32), self.bias.type(torch.float32), self.eps)
return ret.type(orig_type)
class RMSNorm(torch.nn.Module):
def __init__(self, dim: int, eps: float = 1e-6):
super().__init__()
self.eps = eps
self.weight = nn.Parameter(torch.ones(dim))
def _norm(self, x):
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
def forward(self, x):
output = self._norm(x.float()).type_as(x)
return output * self.weight |