| 11 | |
| 12 | |
| 13 | class RMS_norm(nn.Module): |
| 14 | |
| 15 | def __init__(self, dim, channel_first=True, images=True, bias=False): |
| 16 | super().__init__() |
| 17 | broadcastable_dims = (1, 1, 1) if not images else (1, 1) |
| 18 | shape = (dim, *broadcastable_dims) if channel_first else (dim,) |
| 19 | |
| 20 | self.channel_first = channel_first |
| 21 | self.scale = dim**0.5 |
| 22 | self.gamma = nn.Parameter(torch.ones(shape)) |
| 23 | self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0. |
| 24 | |
| 25 | def forward(self, x): |
| 26 | return F.normalize( |
| 27 | x, dim=(1 if self.channel_first else |
| 28 | -1)) * self.scale * self.gamma + self.bias |
| 29 | |
| 30 | class CausalConv3d(nn.Conv3d): |
| 31 | """ |