Initial commit and v1.0

demucs/grids/repro.py

@@ -0,0 +1,50 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+"""
+Easier training for reproducibility
+"""
+
+from ._explorers import MyExplorer
+
+
+@MyExplorer
+def explorer(launcher):
+    launcher.slurm_(
+        gpus=8,
+        time=3 * 24 * 60,
+        partition='devlab,learnlab')
+
+    launcher.bind_({'ema.epoch': [0.9, 0.95]})
+    launcher.bind_({'ema.batch': [0.9995, 0.9999]})
+    launcher.bind_({'epochs': 600})
+
+    base = {'model': 'demucs', 'demucs.dconv_mode': 0, 'demucs.gelu': False,
+            'demucs.lstm_layers': 2}
+    newt = {'model': 'demucs', 'demucs.normalize': True}
+    hdem = {'model': 'hdemucs'}
+    svd = {'svd.penalty': 1e-5, 'svd': 'base2'}
+
+    with launcher.job_array():
+        for model in [base, newt, hdem]:
+            sub = launcher.bind(model)
+            if model is base:
+                # Training the v2 Demucs on MusDB HQ
+                sub(epochs=360)
+                continue
+
+            # those two will be used in the repro_mdx_a bag of models.
+            sub(svd)
+            sub(svd, seed=43)
+            if model == newt:
+                # Ablation study
+                sub()
+                abl = sub.bind(svd)
+                abl({'ema.epoch': [], 'ema.batch': []})
+                abl({'demucs.dconv_lstm': 10})
+                abl({'demucs.dconv_attn': 10})
+                abl({'demucs.dconv_attn': 10, 'demucs.dconv_lstm': 10, 'demucs.lstm_layers': 2})
+                abl({'demucs.dconv_mode': 0})
+                abl({'demucs.gelu': False})