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.
(model: torch.nn.Module, init: str)
| 10 | |
| 11 | @typechecked |
| 12 | def initialize(model: torch.nn.Module, init: str): |
| 13 | """Initialize weights of a neural network module. |
| 14 | |
| 15 | Parameters are initialized using the given method or distribution. |
| 16 | |
| 17 | Custom initialization routines can be implemented into submodules |
| 18 | as function `espnet_initialization_fn` within the custom module. |
| 19 | |
| 20 | Args: |
| 21 | model: Target. |
| 22 | init: Method of initialization. |
| 23 | """ |
| 24 | # weight init |
| 25 | for p in model.parameters(): |
| 26 | if p.dim() > 1: |
| 27 | if init == "xavier_uniform": |
| 28 | torch.nn.init.xavier_uniform_(p.data) |
| 29 | elif init == "xavier_normal": |
| 30 | torch.nn.init.xavier_normal_(p.data) |
| 31 | elif init == "kaiming_uniform": |
| 32 | torch.nn.init.kaiming_uniform_(p.data, nonlinearity="relu") |
| 33 | elif init == "kaiming_normal": |
| 34 | torch.nn.init.kaiming_normal_(p.data, nonlinearity="relu") |
| 35 | elif init == "normal": |
| 36 | torch.nn.init.normal_(p.data, mean=0.0, std=0.02) |
| 37 | else: |
| 38 | raise ValueError("Unknown initialization: " + init) |
| 39 | # bias init |
| 40 | for name, p in model.named_parameters(): |
| 41 | if ".bias" in name and p.dim() == 1: |
| 42 | p.data.zero_() |
| 43 | logging.info(f"Initialize {name} to zeros") |
| 44 | |
| 45 | # reset some modules with default init |
| 46 | for m in model.modules(): |
| 47 | if isinstance(m, (torch.nn.Embedding, torch.nn.LayerNorm, torch.nn.GroupNorm)): |
| 48 | m.reset_parameters() |
| 49 | if hasattr(m, "espnet_initialization_fn"): |
| 50 | m.espnet_initialization_fn() |
| 51 | |
| 52 | # TODO(xkc): Hacking s3prl_frontend and wav2vec2encoder initialization |
| 53 | if getattr(model, "encoder", None) and getattr( |
| 54 | model.encoder, "reload_pretrained_parameters", None |
| 55 | ): |
| 56 | model.encoder.reload_pretrained_parameters() |
| 57 | if getattr(model, "frontend", None): |
| 58 | if getattr(model.frontend, "reload_pretrained_parameters", None): |
| 59 | model.frontend.reload_pretrained_parameters() |
| 60 | elif isinstance( |
| 61 | getattr(model.frontend, "frontends", None), |
| 62 | torch.nn.ModuleList, |
| 63 | ): |
| 64 | for i, _ in enumerate(getattr(model.frontend, "frontends")): |
| 65 | if getattr( |
| 66 | model.frontend.frontends[i], |
| 67 | "reload_pretrained_parameters", |
| 68 | None, |
| 69 | ): |
searching dependent graphs…