File size: 1,652 Bytes
f15d29e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 | """
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
# This function is not type annotated because mypy complains that axis could be either an integer or a tuple of integers,
# even though this is precicely how torch.var_mean works
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) # last dimension is output dimension
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
|