| 849 | self.initialized = not ddi |
| 850 | |
| 851 | def initialize(self, x, x_mask): |
| 852 | with torch.no_grad(): |
| 853 | denom = torch.sum(x_mask, [0, 2]) |
| 854 | m = torch.sum(x * x_mask, [0, 2]) / denom |
| 855 | m_sq = torch.sum(x * x * x_mask, [0, 2]) / denom |
| 856 | v = m_sq - (m**2) |
| 857 | logs = 0.5 * torch.log(torch.clamp_min(v, 1e-6)) |
| 858 | |
| 859 | bias_init = ( |
| 860 | (-m * torch.exp(-logs)).view(*self.bias.shape).to(dtype=self.bias.dtype) |
| 861 | ) |
| 862 | logs_init = (-logs).view(*self.logs.shape).to(dtype=self.logs.dtype) |
| 863 | |
| 864 | self.bias.data.copy_(bias_init) |
| 865 | self.logs.data.copy_(logs_init) |
| 866 | |
| 867 | |
| 868 | class InvConvNear(nn.Module): |