fix and optimize

Author: yxlllc

configs/acoustic.yaml

@@ -70,7 +70,7 @@ sampling_algorithm: euler
 sampling_steps: 20
 diff_accelerator: ddim
 diff_speedup: 10
-hidden_size: 256
+hidden_size: 384
 backbone_type: 'lynxnet2'
 backbone_args:
   num_channels: 1024

configs/templates/config_acoustic.yaml

@@ -77,6 +77,7 @@ T_start: 0.4
 T_start_infer: 0.4
 K_step: 300
 K_step_infer: 300
+hidden_size: 384
 backbone_type: 'lynxnet2'
 backbone_args:
   num_channels: 1024

modules/backbones/lynxnet.py

@@ -74,7 +74,7 @@ def forward(self, x, conditioner, diffusion_step, front_cond_inject=False):
 
 class LYNXNet(nn.Module):
     def __init__(self, in_dims, n_feats, *, num_layers=6, num_channels=512, expansion_factor=2, kernel_size=31,
-                 activation='PReLU', dropout=0.0, strong_cond=False):
+                 activation='PReLU', dropout_rate=0.0, strong_cond=False):
         """
         LYNXNet(Linear Gated Depthwise Separable Convolution Network)
         TIPS:You can control the style of the generated results by modifying the 'activation', 
@@ -100,7 +100,7 @@ def __init__(self, in_dims, n_feats, *, num_layers=6, num_channels=512, expansio
                     expansion_factor=expansion_factor,
                     kernel_size=kernel_size,
                     activation=activation,
-                    dropout=dropout
+                    dropout=dropout_rate
                 )
                 for i in range(num_layers)
             ]

modules/backbones/lynxnet2.py

@@ -39,7 +39,7 @@ def forward(self, x):
 
 class LYNXNet2(nn.Module):
     def __init__(self, in_dims, n_feats, *, num_layers=6, num_channels=512, expansion_factor=1, kernel_size=31,
-                 dropout=0.0, use_conditioner_cache=False, glu_type='swiglu'):
+                 dropout_rate=0.0, use_conditioner_cache=False, glu_type='swiglu'):
         """
         LYNXNet2(Linear Gated Depthwise Separable Convolution Network Version 2)
         """
@@ -65,7 +65,7 @@ def __init__(self, in_dims, n_feats, *, num_layers=6, num_channels=512, expansio
                     dim=num_channels,
                     expansion_factor=expansion_factor,
                     kernel_size=kernel_size,
-                    dropout=dropout,
+                    dropout=dropout_rate,
                     glu_type=glu_type
                 )
                 for i in range(num_layers)

modules/optimizer/muon.py

@@ -37,25 +37,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 G.size(-2) > G.size(-1):
-        X = X.mT
-    return X.to(G)
+    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)
+            
+    return X
 
 
 class Muon(torch.optim.Optimizer):
@@ -85,7 +85,7 @@ def __init__(self, params, lr=5e-4, weight_decay=0.1, momentum=0.95, nesterov=Tr
         defaults = dict(lr=lr, weight_decay=weight_decay, momentum=momentum, nesterov=nesterov, ns_steps=ns_steps)
         super().__init__(params, defaults)
         self.bf16_support_map = get_bf16_support_map()
-        
+    
     @torch.no_grad()
     def step(self, closure=None):
         for group in self.param_groups:
@@ -95,28 +95,29 @@ def step(self, closure=None):
                 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]: