Upload luminars/train.py

luminars/train.py

@@ -0,0 +1,26 @@
+"""Staged freeze/thaw training helpers."""
+import torch.nn as nn
+
+def freeze(m):
+    for p in m.parameters(): p.requires_grad = False
+
+def unfreeze(m):
+    for p in m.parameters(): p.requires_grad = True
+
+def count_trainable(m):
+    return sum(p.numel() for p in m.parameters() if p.requires_grad)
+
+def setup_stage1(model):
+    unfreeze(model)
+    return model
+
+def setup_stage2(model):
+    freeze(model)
+    for name, p in model.named_parameters():
+        if "xattn" in name or "tproj" in name:
+            p.requires_grad = True
+    return model
+
+def setup_stage3(model):
+    unfreeze(model)
+    return model