Layer scaling module for stable training.
| 5 | |
| 6 | |
| 7 | class LayerScale(nn.Module): |
| 8 | """Layer scaling module for stable training.""" |
| 9 | |
| 10 | def __init__( |
| 11 | self, |
| 12 | dim: int, |
| 13 | init_values: Union[float, Tensor] = 1e-5, |
| 14 | inplace: bool = False, |
| 15 | device=None, |
| 16 | ) -> None: |
| 17 | super().__init__() |
| 18 | self.inplace = inplace |
| 19 | self.gamma = nn.Parameter(torch.empty(dim, device=device)) |
| 20 | self.init_values = init_values |
| 21 | |
| 22 | def reset_parameters(self): |
| 23 | nn.init.constant_(self.gamma, self.init_values) |
| 24 | |
| 25 | def forward(self, x: Tensor) -> Tensor: |
| 26 | return x.mul_(self.gamma) if self.inplace else x * self.gamma |
| 27 | |
| 28 | |
| 29 | class PatchDropout(nn.Module): |