| """ |
| Copyright (c) Facebook, Inc. and its affiliates. |
| Copyright (c) Microsoft Corporation. |
| Licensed under the MIT License. |
| Adapted from https://github.com/FAIR-Chem/fairchem/blob/main/src/fairchem/core/models/gemnet/initializers.py. |
| """ |
|
|
| import torch |
|
|
|
|
| |
| |
| def _standardize(kernel): |
| """ |
| Makes sure that N*Var(W) = 1 and E[W] = 0 |
| """ |
| eps = 1e-6 |
|
|
| if len(kernel.shape) == 3: |
| axis = (0, 1) |
| else: |
| axis = 1 |
|
|
| var, mean = torch.var_mean(kernel, dim=axis, unbiased=True, keepdim=True) |
| kernel = (kernel - mean) / (var + eps) ** 0.5 |
| return kernel |
|
|
|
|
| def he_orthogonal_init(tensor: torch.Tensor) -> torch.Tensor: |
| """ |
| Generate a weight matrix with variance according to He (Kaiming) initialization. |
| Based on a random (semi-)orthogonal matrix neural networks |
| are expected to learn better when features are decorrelated |
| (stated by eg. "Reducing overfitting in deep networks by decorrelating representations", |
| "Dropout: a simple way to prevent neural networks from overfitting", |
| "Exact solutions to the nonlinear dynamics of learning in deep linear neural networks") |
| """ |
| tensor = torch.nn.init.orthogonal_(tensor) |
|
|
| if len(tensor.shape) == 3: |
| fan_in = tensor.shape[:-1].numel() |
| else: |
| fan_in = tensor.shape[1] |
|
|
| with torch.no_grad(): |
| tensor.data = _standardize(tensor.data) |
| tensor.data *= (1 / fan_in) ** 0.5 |
|
|
| return tensor |
|
|
