File size: 3,321 Bytes
ad16788 |
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 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 |
#!/usr/bin/env python3
"""Initialize modules for espnet2 neural networks."""
import math
import torch
from typeguard import check_argument_types
def initialize(model: torch.nn.Module, init: str):
"""Initialize weights of a neural network module.
Parameters are initialized using the given method or distribution.
Custom initialization routines can be implemented into submodules
as function `espnet_initialization_fn` within the custom module.
Args:
model: Target.
init: Method of initialization.
"""
assert check_argument_types()
if init == "chainer":
# 1. lecun_normal_init_parameters
for p in model.parameters():
data = p.data
if data.dim() == 1:
# bias
data.zero_()
elif data.dim() == 2:
# linear weight
n = data.size(1)
stdv = 1.0 / math.sqrt(n)
data.normal_(0, stdv)
elif data.dim() in (3, 4):
# conv weight
n = data.size(1)
for k in data.size()[2:]:
n *= k
stdv = 1.0 / math.sqrt(n)
data.normal_(0, stdv)
else:
raise NotImplementedError
for mod in model.modules():
# 2. embed weight ~ Normal(0, 1)
if isinstance(mod, torch.nn.Embedding):
mod.weight.data.normal_(0, 1)
# 3. forget-bias = 1.0
elif isinstance(mod, torch.nn.RNNCellBase):
n = mod.bias_ih.size(0)
mod.bias_ih.data[n // 4 : n // 2].fill_(1.0)
elif isinstance(mod, torch.nn.RNNBase):
for name, param in mod.named_parameters():
if "bias" in name:
n = param.size(0)
param.data[n // 4 : n // 2].fill_(1.0)
if hasattr(mod, "espnet_initialization_fn"):
mod.espnet_initialization_fn()
else:
# weight init
for p in model.parameters():
if p.dim() > 1:
if init == "xavier_uniform":
torch.nn.init.xavier_uniform_(p.data)
elif init == "xavier_normal":
torch.nn.init.xavier_normal_(p.data)
elif init == "kaiming_uniform":
torch.nn.init.kaiming_uniform_(p.data, nonlinearity="relu")
elif init == "kaiming_normal":
torch.nn.init.kaiming_normal_(p.data, nonlinearity="relu")
else:
raise ValueError("Unknown initialization: " + init)
# bias init
for p in model.parameters():
if p.dim() == 1:
p.data.zero_()
# reset some modules with default init
for m in model.modules():
if isinstance(m, (torch.nn.Embedding, torch.nn.LayerNorm)):
m.reset_parameters()
if hasattr(m, "espnet_initialization_fn"):
m.espnet_initialization_fn()
# TODO(xkc): Hacking wav2vec2 initialization
if getattr(model, "encoder", None) and getattr(
model.encoder, "reload_pretrained_parameters", None
):
model.encoder.reload_pretrained_parameters()
|