| 54 | |
| 55 | |
| 56 | class RMS_norm(nn.Module): |
| 57 | |
| 58 | def __init__(self, dim, channel_first=True, images=True, bias=False): |
| 59 | super().__init__() |
| 60 | broadcastable_dims = (1, 1, 1) if not images else (1, 1) |
| 61 | shape = (dim, *broadcastable_dims) if channel_first else (dim,) |
| 62 | |
| 63 | self.channel_first = channel_first |
| 64 | self.scale = dim**0.5 |
| 65 | self.gamma = nn.Parameter(torch.ones(shape)) |
| 66 | self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0. |
| 67 | |
| 68 | def forward(self, x): |
| 69 | return F.normalize( |
| 70 | x, dim=(1 if self.channel_first else |
| 71 | -1)) * self.scale * self.gamma + self.bias |
| 72 | |
| 73 | |
| 74 | class Upsample(nn.Upsample): |