Initial commit: FunASR Speech Recognition Toolkit

Add complete FunASR codebase including models, runtime, and documentation.

funasr/train_utils/initialize.py

@@ -0,0 +1,55 @@
+#!/usr/bin/env python3
+
+"""Initialize modules for espnet2 neural networks."""
+
+import math
+import torch
+
+
+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.
+    """
+
+    # 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, torch.nn.GroupNorm)):
+            m.reset_parameters()
+        if hasattr(m, "espnet_initialization_fn"):
+            m.espnet_initialization_fn()
+
+    # TODO(xkc): Hacking s3prl_frontend and wav2vec2encoder initialization
+    if getattr(model, "encoder", None) and getattr(
+        model.encoder, "reload_pretrained_parameters", None
+    ):
+        model.encoder.reload_pretrained_parameters()
+    if getattr(model, "frontend", None) and getattr(
+        model.frontend, "reload_pretrained_parameters", None
+    ):
+        model.frontend.reload_pretrained_parameters()