full pipeline

init_cross_attn.py

@@ -9,6 +9,17 @@ from lightning.pytorch import seed_everything
 import random
 from datetime import datetime
 import numpy as np
+import sys
+
+# 获取当前脚本所在目录（ckpts/）
+current_dir = os.path.dirname(os.path.abspath(__file__))
+
+# 项目根目录 = ckpts 的上级目录
+project_root = os.path.abspath(os.path.join(current_dir, '..'))
+
+# 添加项目根目录到 sys.path
+if project_root not in sys.path:
+    sys.path.insert(0, project_root)
 
 from ThinkSound.data.datamodule import DataModule
 from ThinkSound.models import create_model_from_config
@@ -62,6 +73,7 @@ def main():
 
     # Step 5: 初始化 cross-attn 模块（只初始化新增部分）
     def init_cross_attn_weights(module):
+        from einops.layers.torch import Rearrange
         if isinstance(module, nn.Linear):
             nn.init.xavier_uniform_(module.weight)
             if module.bias is not None:
@@ -69,7 +81,14 @@ def main():
         elif isinstance(module, nn.LayerNorm):
             nn.init.ones_(module.weight)
             nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.RMSNorm) or module.__class__.__name__ == "RMSNorm":
+            if hasattr(module, 'weight'):
+                nn.init.ones_(module.weight)
+            if hasattr(module, 'bias') and module.bias is not None:
+                nn.init.zeros_(module.bias)
 
+    import pdb; pdb.set_trace()
+    pass
     # 只遍历 cross-attn 模块进行初始化
     for name, module in model.named_modules():
         if 'cross_attn' in name:
@@ -77,7 +96,7 @@ def main():
             print(f"[INIT] Initialized {name}")
 
     # Step 6: 保存新权重
-    torch.save(model.state_dict(), 'ckpts/thinksound_light_cross_attn.ckpt')
+    torch.save(model.state_dict(), 'ckpts/row_thinksound_light_cross_attn.ckpt')
     print("[DONE] New checkpoint saved with old weights + initialized cross-attn.")
 
 if __name__ == '__main__':

row_thinksound_light_cross_attn.ckpt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1a80397e1a4f44c9e698ce2b6cbacf5cc775ef598907e41b2241b285b9e7eb78
-size 5909451670
+oid sha256:db9f641e234f91448d9ca1ec9254339ad64414cbdc2c637311ff06b985d8fb65
+size 6026895638

thinksound_light_cross_attn.ckpt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:1a80397e1a4f44c9e698ce2b6cbacf5cc775ef598907e41b2241b285b9e7eb78
-size 5909451670