house cleaning

Author: wkzng

datasets/compiler.py

@@ -167,23 +167,20 @@ def run(self, max_samples:int=None):
 
 
 if __name__ == "__main__":
+    ROOT = os.getenv("ROOT_DIRECTORY")
     SAMPLE_RATE = int(os.getenv("SAMPLE_RATE"))
     CHUNK_DURATION = int(os.getenv("CHUNK_DURATION"))
     HOP_LENGTH = int(os.getenv("HOP_LENGTH"))
     N_SAMPLES = int(SAMPLE_RATE * CHUNK_DURATION)
+    TARGET_LUFFS = -float(os.getenv("TARGET_LUFFS", 23))
     N_CHANNELS = 1
-    TARGET_LUFFS = -23
-    ROOT = os.getenv("ROOT_DIRECTORY")
 
     buffer_schemas = {
         "audio": (N_CHANNELS, N_SAMPLES), 
     }
 
     datasets_folders = [
         'gtzan/train',
-        'mvsep_multisong_dataset/train',
-        'mvsep_synth_dataset/train',
-        'djmax_respectv_22050/train',
     ]
 
     datasets_config = [
@@ -198,7 +195,7 @@ def run(self, max_samples:int=None):
         for folder in datasets_folders
     ]
 
-    compile_name = "mixture"
+    compile_name = "gtzan"
     compile_dir = os.path.join(ROOT, "_compiled")
     os.makedirs(compile_dir, exist_ok=True)
 

demo.ipynb

@@ -126,13 +126,13 @@ def train(self):
 
 
 if __name__ =="__main__":
-    from datasets.configs import GTZANConfig
+    from datasets.configs import BaseDatasetConfig
     from datasets.dataset import ChunkDataset
     from sincnet.model import SincNet
 
-    model = SincNet()
-    dataset_config = GTZANConfig(id="gtzan")
-    learning_rate = 1e-3
+    model = SincNet(scale="mel")
+    dataset_config = BaseDatasetConfig(id="gtzan")
+    learning_rate = 1e-4
     train_config = TrainConfig(**{
         "batch_size": 8,
         "n_epoch": 500,
@@ -157,4 +157,10 @@ def train(self):
         val_set=datasets["test"],
         config=train_config
     )
+
+    try:
+        model.load_pretrained_weights(weights_folder="pretrained", freeze=False)
+    except:
+        pass
+
     trainer.train()