| |
| """Initialize modules for espnet2 neural networks.""" |
| 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() |
| print("init with", 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) |
| |
| for name, p in model.named_parameters(): |
| if ".bias" in name and p.dim() == 1: |
| p.data.zero_() |
|
|
