Merge pull request #9 from wolfgitpr/aug

Aug

Author: wolfgitpr

binarize.py

@@ -1,4 +1,4 @@
-model_name: 1201_fa
+model_name: 1218_fa
 
 float32_matmul_precision: high
 random_seed: 123456
@@ -67,7 +67,7 @@ fa_arg:
 optimizer_config:
   lr: 0.0003
   gamma: 0.9999
-  total_steps: 20000
+  total_steps: 10000
   muon_args:
     weight_decay: 0.1
   adamw_args:
@@ -104,9 +104,9 @@ hubert_config:
 mel_spec_config:
   n_mels: 128
   sample_rate: 44100
-  window_size: 1024
-  hop_size: 512
-  n_fft: 2048
+  window_size: 882
+  hop_size: 441
+  n_fft: 1764
   f_min: 40
   f_max: 16000
   clamp: 0.00001

configs/force_alignment.yaml

@@ -1,4 +1,4 @@
-model_name: 1201_nll
+model_name: 1218_nll
 
 float32_matmul_precision: high
 random_seed: 123456
@@ -54,7 +54,7 @@ cvnt_arg:
 optimizer_config:
   lr: 0.0003
   gamma: 0.9999
-  total_steps: 20000
+  total_steps: 10000
   muon_args:
     weight_decay: 0.1
   adamw_args:
@@ -86,9 +86,9 @@ hubert_config:
 mel_spec_config:
   n_mels: 128
   sample_rate: 44100
-  window_size: 1024
-  hop_size: 512
-  n_fft: 2048
+  window_size: 882
+  hop_size: 441
+  n_fft: 1764
   f_min: 40
   f_max: 16000
   clamp: 0.00001

configs/non_lexical_labeler.yaml

@@ -0,0 +1,19 @@
+import torch.nn as nn
+
+
+class SEAttention(nn.Module):
+    def __init__(self, channels, reduction=16):
+        super().__init__()
+        self.fc = nn.Sequential(
+            nn.AdaptiveAvgPool1d(1),
+            nn.Flatten(),
+            nn.Linear(channels, channels // reduction, bias=False),
+            nn.ReLU(inplace=True),
+            nn.Linear(channels // reduction, channels, bias=False),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):  # x: [B, C, T]
+        B, C, T = x.shape
+        scale = self.fc(x).view(B, C, 1)  # [B, C, 1]
+        return x * scale.expand_as(x)

networks/layer/attention/__init__.py

@@ -1,6 +1,8 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
+from networks.layer.attention.attention import SEAttention
+
 
 class UNetBackbone(nn.Module):
     def __init__(
@@ -16,26 +18,15 @@ def __init__(
             channels_scaleup_factor: float = 2.0,
             dropout: float = 0.1,
     ):
-        """_summary_
-
-        Args:
-            input_dims (int):
-            output_dims (int):
-            hidden_dims (int):
-            block (nn.Module): shape: (B, C, T) -> shape: (B, C, T)
-            down_sampling (nn.Module): shape: (B, C, T) -> shape: (B, C*2, T/down_sampling_factor)
-            up_sampling (nn.Module): shape: (B, C, T) -> shape: (B, C/2, T*down_sampling_factor)
-        """
         super(UNetBackbone, self).__init__()
 
         self.input_dims = input_dims
         self.output_dims = output_dims
         self.hidden_dims = hidden_dims
-
         self.divisible_factor = down_sampling_factor ** down_sampling_times
 
         self.encoders = nn.ModuleList()
-        self.encoders.append(block(input_dims, hidden_dims, dropout=dropout))
+        self.encoders.append(self._attention_block(block, input_dims, hidden_dims, dropout))
         for i in range(down_sampling_times - 1):
             i += 1
             self.encoders.append(
@@ -45,10 +36,11 @@ def __init__(
                         int(channels_scaleup_factor ** i) * hidden_dims,
                         down_sampling_factor,
                     ),
-                    block(
+                    self._attention_block(
+                        block,
                         int(channels_scaleup_factor ** i) * hidden_dims,
                         int(channels_scaleup_factor ** i) * hidden_dims,
-                        dropout=dropout,
+                        dropout,
                     ),
                 )
             )
@@ -59,10 +51,11 @@ def __init__(
                 int(channels_scaleup_factor ** down_sampling_times) * hidden_dims,
                 down_sampling_factor,
             ),
-            block(
+            self._attention_block(
+                block,
                 int(channels_scaleup_factor ** down_sampling_times) * hidden_dims,
                 int(channels_scaleup_factor ** down_sampling_times) * hidden_dims,
-                dropout=dropout,
+                dropout,
             ),
             up_sampling(
                 int(channels_scaleup_factor ** down_sampling_times) * hidden_dims,
@@ -76,12 +69,13 @@ def __init__(
             i += 1
             self.decoders.append(
                 nn.Sequential(
-                    block(
+                    self._attention_block(
+                        block,
                         int(channels_scaleup_factor ** (down_sampling_times - i))
                         * hidden_dims,
                         int(channels_scaleup_factor ** (down_sampling_times - i))
                         * hidden_dims,
-                        dropout=dropout,
+                        dropout,
                     ),
                     up_sampling(
                         int(channels_scaleup_factor ** (down_sampling_times - i))
@@ -92,11 +86,16 @@ def __init__(
                     ),
                 )
             )
-        self.decoders.append(block(hidden_dims, output_dims))
+        self.decoders.append(self._attention_block(block, hidden_dims, output_dims, dropout))
+
+    @staticmethod
+    def _attention_block(block_class, in_channels, out_channels, dropout):
+        return nn.Sequential(
+            block_class(in_channels, out_channels, dropout=dropout),
+            SEAttention(out_channels)
+        )
 
-    def forward(self,
-                x,  # [B, C, T]
-                ):
+    def forward(self, x):  # x: [B, C, T]
         B, C, T = x.shape
         padding_len = (-T) % self.divisible_factor
         x = F.pad(x, (0, padding_len))

networks/layer/attention/attention.py

@@ -1,5 +1,7 @@
 import torch.nn as nn
 
+from networks.layer.attention.attention import SEAttention
+
 
 class ResidualBasicBlock(nn.Module):
     def __init__(self, input_dims: int, output_dims: int, hidden_dims: int = None, n_groups: int = 16,
@@ -37,6 +39,47 @@ def forward(self, x):  # x: [B, C, T]
         return x
 
 
+class AttentionResidualBasicBlock(nn.Module):
+    def __init__(self, input_dims: int, output_dims: int, hidden_dims: int = None,
+                 n_groups: int = 16, dropout: float = 0.1):
+        super().__init__()
+
+        self.input_dims = input_dims
+        self.output_dims = output_dims
+        self.hidden_dims = hidden_dims if hidden_dims is not None else max(n_groups, output_dims // n_groups * n_groups)
+        self.n_groups = n_groups
+        self.dropout = dropout
+
+        self.activation = nn.Hardswish()
+
+        self.block = nn.Sequential(
+            nn.Conv1d(self.input_dims, self.hidden_dims, kernel_size=3, padding=1, bias=False),
+            nn.GroupNorm(self.n_groups, self.hidden_dims),
+            self.activation,
+            nn.Dropout1d(dropout) if dropout > 0 else nn.Identity(),
+            nn.Conv1d(self.hidden_dims, self.output_dims, kernel_size=3, padding=1, bias=False),
+        )
+
+        self.attention = SEAttention(self.output_dims)
+
+        self.shortcut = nn.Identity() if self.input_dims == self.output_dims \
+            else nn.Conv1d(self.input_dims, self.output_dims, kernel_size=1, bias=False)
+
+        self.out = nn.Sequential(
+            nn.GroupNorm(1, self.output_dims),
+            self.activation,
+            nn.Dropout(dropout) if dropout > 0 else nn.Identity(),
+        )
+
+    def forward(self, x):
+        residual = self.shortcut(x)
+        x = self.block(x)
+        x = self.attention(x)
+        x = x + residual
+        x = self.out(x)
+        return x
+
+
 class ResidualBottleNeckBlock(nn.Module):
     def __init__(self, input_dims: int, output_dims: int, hidden_dims: int = None, n_groups: int = 16):
         super(ResidualBottleNeckBlock, self).__init__()

networks/layer/backbone/unet.py

@@ -39,25 +39,25 @@ def zeropower_via_newtonschulz5(G: Tensor, steps: int, use_bf16: bool) -> Tensor
     """
     assert G.ndim == 3  # batched Muon implementation by @scottjmaddox, and put into practice in the record by @YouJiacheng
     a, b, c = (3.4445, -4.7750, 2.0315)
-    if use_bf16:
-        X = G.bfloat16()
-    else:
-        X = G.float()
-    if G.size(-2) > G.size(-1):
-        X = X.mT
+
+    X = G.to(dtype=torch.bfloat16 if use_bf16 else torch.float32)
 
     # Ensure spectral norm is at most 1
     X = F.normalize(X, p=2.0, dim=(-2, -1), eps=1e-7)
 
     # Perform the NS iterations
-    for _ in range(steps):
-        A = X @ X.mT
-        B = torch.baddbmm(A, A, A, beta=b, alpha=c)
-        X = torch.baddbmm(X, B, X, beta=a, alpha=1)
+    if X.size(-2) < X.size(-1):
+        for _ in range(steps):
+            A = torch.bmm(X, X.mT)
+            A = torch.baddbmm(A, A, A, beta=b, alpha=c)
+            X = torch.baddbmm(X, A, X, beta=a, alpha=1)
+    else:
+        for _ in range(steps):
+            A = torch.bmm(X.mT, X)
+            A = torch.baddbmm(A, A, A, beta=b, alpha=c)
+            X = torch.baddbmm(X, X, A, beta=a, alpha=1)
 
-    if G.size(-2) > G.size(-1):
-        X = X.mT
-    return X.to(G)
+    return X
 
 
 class Muon(torch.optim.Optimizer):
@@ -92,33 +92,34 @@ def __init__(self, params, lr=5e-4, weight_decay=0.1, momentum=0.95, nesterov=Tr
     def step(self, closure=None):
         for group in self.param_groups:
             shape_groups = {}
-            for p in filter(lambda p: p.grad is not None, group["params"]):
+            for p in filter(lambda _p: _p.grad is not None, group["params"]):
                 g = p.grad
                 state = self.state[p]
                 if "momentum_buffer" not in state:
                     state["momentum_buffer"] = torch.zeros_like(g)
-                buf: Tensor = state["momentum_buffer"]
                 key = (p.shape, p.device, p.dtype)
                 if key not in shape_groups:
                     shape_groups[key] = {"params": [], "grads": [], "buffers": []}
                 shape_groups[key]["params"].append(p)
                 shape_groups[key]["grads"].append(g)
-                shape_groups[key]["buffers"].append(buf)
+                shape_groups[key]["buffers"].append(state["momentum_buffer"])
             for key in shape_groups:
                 group_data = shape_groups[key]
-                g = torch.stack(group_data["grads"])
-                buf = torch.stack(group_data["buffers"])
-                buf.lerp_(g, 1 - group["momentum"])
-                g = g.lerp_(buf, group["momentum"]) if group["nesterov"] else buf
+                p, g, buf, m = group_data["params"], group_data["grads"], group_data["buffers"], group["momentum"]
+                torch._foreach_lerp_(buf, g, 1 - m)
+                if group["nesterov"]:
+                    torch._foreach_lerp_(g, buf, m)
+                    g = torch.stack(g)
+                else:
+                    g = torch.stack(buf)
+                original_shape = g.shape
                 if g.ndim >= 4:  # for the case of conv filters
                     g = g.view(g.size(0), g.size(1), -1)
                 use_bf16 = self.bf16_support_map.get(g.device, False)
                 g = zeropower_via_newtonschulz5(g, steps=group["ns_steps"], use_bf16=use_bf16)
-                for i, p in enumerate(group_data["params"]):
-                    if group["weight_decay"] > 0:
-                        p.data.mul_(1 - group["lr"] * group["weight_decay"])
-                    p.data.add_(g[i].view_as(p), alpha=-group["lr"] * max(g[i].size()) ** 0.5)
-                    self.state[p]["momentum_buffer"] = buf[i].clone()
+                if group["weight_decay"] > 0:
+                    torch._foreach_mul_(p, 1 - group["lr"] * group["weight_decay"])
+                torch._foreach_add_(p, g.view(original_shape).unbind(0), alpha=-group["lr"] * max(g[0].size()) ** 0.5)
 
 
 def get_params_for_muon(model) -> List[Parameter]:
@@ -129,6 +130,7 @@ def get_params_for_muon(model) -> List[Parameter]:
         module: The module to filter parameters for.
     Returns:
         A list of parameters that should be optimized with muon.
+        :param model:
     """
     muon_params = []
     for module in model.modules():
@@ -141,7 +143,11 @@ def get_params_for_muon(model) -> List[Parameter]:
 
 
 class Muon_AdamW(ChainedOptimizer):
-    def __init__(self, model, lr=0.0005, weight_decay=0.0, muon_args={}, adamw_args={}, verbose=False):
+    def __init__(self, model, lr=0.0005, weight_decay=0.0, muon_args=None, adamw_args=None, verbose=False):
+        if adamw_args is None:
+            adamw_args = {}
+        if muon_args is None:
+            muon_args = {}
         muon_params_id_set = set(id(p) for p in get_params_for_muon(model))
         spec_muon = OptimizerSpec(Muon, muon_args, lambda param: id(param) in muon_params_id_set)
         spec_adamw = OptimizerSpec(torch.optim.AdamW, adamw_args, None)

networks/layer/block/resnet_block.py

@@ -6,15 +6,14 @@
 import torch.optim.lr_scheduler as lr_scheduler_module
 
 import networks.scheduler as scheduler_module
-from scripts.evaluate import remove_ignored_phonemes, quantize_tier
 from networks.layer.backbone.unet import UNetBackbone
-from networks.layer.block.resnet_block import ResidualBasicBlock
+from networks.layer.block.resnet_block import AttentionResidualBasicBlock
 from networks.layer.fusion.curves_fusion import PowerCurveEdgeFusion
 from networks.layer.scaling.stride_conv import DownSampling, UpSampling
 from networks.loss.GHMLoss import CTCGHMLoss, GHMLoss, MultiLabelGHMLoss
 from networks.optimizer.muon import Muon_AdamW
+from scripts.evaluate import remove_ignored_phonemes, quantize_tier
 from tools.decoder import AlignmentDecoder
-from tools.encoder import UnitsEncoder
 from tools.metrics import BoundaryEditRatio, BoundaryEditRatioWeighted, VlabelerEditRatio, CustomPointTier
 
 
@@ -49,12 +48,12 @@ def __init__(
             input_dims=self.hubert_config["channel"],
             output_dims=self.fa_arg["hidden_dims"],
             hidden_dims=self.fa_arg["hidden_dims"],
-            block=ResidualBasicBlock,
+            block=AttentionResidualBasicBlock,
             down_sampling=DownSampling,
             up_sampling=UpSampling,
-            down_sampling_factor=self.fa_arg["down_sampling_factor"],  # 2
-            down_sampling_times=self.fa_arg["down_sampling_times"],  # 3
-            channels_scaleup_factor=self.fa_arg["channels_scaleup_factor"],  # 1.5
+            down_sampling_factor=self.fa_arg["down_sampling_factor"],
+            down_sampling_times=self.fa_arg["down_sampling_times"],
+            channels_scaleup_factor=self.fa_arg["channels_scaleup_factor"],
             dropout=self.fa_arg["dropout"],
         )
 
@@ -145,10 +144,6 @@ def _losses_schedulers_step(self):
     def _losses_schedulers_call(self):
         return torch.tensor([scheduler() for scheduler in self.losses_schedulers]).to(self.device)
 
-    def on_predict_start(self):
-        if self.unitsEncoder is None:
-            self.unitsEncoder = UnitsEncoder(self.hubert_config, self.mel_spec_config, device=self.device)
-
     def _get_consistency_loss(
             self,
             ph_frame_logits,  # [B,vocab_size,T]