Fix some issue about Initialization (#250)

issue 249(2/3/4)

Author: KakaruHayate

modules/backbones/lynxnet.py

@@ -6,7 +6,8 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-from modules.commons.common_layers import SinusoidalPosEmb, SwiGLU, Conv1d, Transpose
+from modules.commons.common_layers import SinusoidalPosEmb, SwiGLU, Transpose
+from modules.commons.common_layers import KaimingNormalConv1d as Conv1d
 from utils.hparams import hparams
 
 

modules/backbones/lynxnet2.py

@@ -2,7 +2,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-from modules.commons.common_layers import SinusoidalPosEmb, SwiGLU, Conv1d, Transpose
+from modules.commons.common_layers import SinusoidalPosEmb, SwiGLU, Transpose
 from utils.hparams import hparams
 
 
@@ -42,6 +42,8 @@ def __init__(self, in_dims, n_feats, *, num_layers=6, num_channels=512, expansio
         self.n_feats = n_feats
         self.input_projection = nn.Linear(in_dims * n_feats, num_channels)
         self.conditioner_projection = nn.Linear(hparams['hidden_size'], num_channels)
+        # It may need to be modified at some point to be compatible with the condition cache
+        # self.conditioner_projection = nn.Conv1d(hparams['hidden_size'], num_channels, 1)
         self.diffusion_embedding = nn.Sequential(
             SinusoidalPosEmb(num_channels),
             nn.Linear(num_channels, num_channels * 4),
@@ -80,6 +82,8 @@ def forward(self, spec, diffusion_step, cond):
 
         x = self.input_projection(x.transpose(1, 2)) # [B, T, F x M]
         x = x + self.conditioner_projection(cond.transpose(1, 2))
+        # It may need to be modified at some point to be compatible with the condition cache
+        # x = x + self.conditioner_projection(cond.transpose(1, 2))
         x = x + self.diffusion_embedding(diffusion_step).unsqueeze(1)
 
         for layer in self.residual_layers:

modules/backbones/wavenet.py

@@ -5,7 +5,8 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-from modules.commons.common_layers import SinusoidalPosEmb, Conv1d
+from modules.commons.common_layers import SinusoidalPosEmb
+from modules.commons.common_layers import KaimingNormalConv1d as Conv1d
 from utils.hparams import hparams
 
 

modules/commons/common_layers.py

@@ -121,11 +121,11 @@ def forward(self, x):
             max_abs_out = torch.max(-out_min, out_max).float()
             max_abs_gate = torch.max(-gate_min, gate_max).float()
             if max_abs_out * max_abs_gate > 1000:
-                return (out.float() * gate.float()).clamp(-1000, 1000).half()             
+                return (out.float() * gate.float()).clamp(-1000, 1000).half()
         return out * gate
 
 
-class Conv1d(torch.nn.Conv1d):
+class KaimingNormalConv1d(torch.nn.Conv1d):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         nn.init.kaiming_normal_(self.weight)
@@ -190,10 +190,17 @@ def __init__(self, embed_dim, num_heads, dropout=0.1, bias=False, rotary_embed=N
 
         # Dropout layer
         self.dropout = nn.Dropout(dropout)
-
+        
         # Rotary Embeddings
         self.rotary_embed = rotary_embed
 
+        # Initialization parameters
+        nn.init.xavier_uniform_(self.in_proj.weight)
+        nn.init.xavier_uniform_(self.out_proj.weight)
+        if bias:
+            nn.init.constant_(self.in_proj.bias, 0.0)
+            nn.init.constant_(self.out_proj.bias, 0.0)
+        
     def forward(self, x, key_padding_mask=None):
         # x: (B, L, C)
         # key_padding_mask: (B, L)