working commit

main.py

@@ -3,34 +3,38 @@ if __name__ == '__main__':
     from tqdm import tqdm
     from cfg import *
     from mimii import MIMII
-    from models.ae import AE
+    from models.ae import AE, SubSpecCAE
     import torch.nn as nn
     import torch.optim as optim
     import random
+    from models.layers import Subspectrogram
 
     torch.manual_seed(42)
     torch.cuda.manual_seed(42)
     np.random.seed(42)
     random.seed(42)
 
-    dataset_path = ALL_DATASET_PATHS[0]
+    dataset_path = ALL_DATASET_PATHS[5]
     print(f'Training on {dataset_path.name}')
     mimii = MIMII(dataset_path=dataset_path, machine_id=0)
     mimii.to(DEVICE)
-    #mimii.preprocess(n_fft=1024, hop_length=256, n_mels=80, center=False, power=2.0)
+    #mimii.preprocess(n_fft=1024, hop_length=256, n_mels=80, center=False, power=2.0)  # 80 x 80
+    tfms = Subspectrogram(SUB_SPEC_HEIGT, SUB_SPEC_HOP)
 
     dl = mimii.train_dataloader(
         segment_len=NUM_SEGMENTS,
         hop_len=NUM_SEGMENT_HOPS,
         batch_size=BATCH_SIZE,
         num_workers=NUM_WORKERS,
-        shuffle=True
+        shuffle=True,
+        transform=tfms
     )
 
-
-    model = AE(400).to(DEVICE)
+    model = SubSpecCAE().to(DEVICE)
     model.init_weights()
 
+    # print(model(torch.randn(128, 1, 20, 80).to(DEVICE)).shape)
+
     optimizer = optim.Adam(model.parameters(), lr=0.001)
 
 
@@ -39,7 +43,7 @@ if __name__ == '__main__':
         losses = []
         for batch in tqdm(dl):
             data, labels = batch
-            data = data.to(DEVICE)
+            data = data.to(DEVICE)  # torch.Size([128, 4, 20, 80]) batch x subs_specs x height x width
 
             loss = model.train_loss(data)
 
@@ -50,7 +54,7 @@ if __name__ == '__main__':
             losses.append(loss.item())
         print(f'Loss: {np.mean(losses)}')
 
-    auc = mimii.evaluate_model(model, NUM_SEGMENTS, NUM_SEGMENTS)
+    auc = mimii.evaluate_model(model, NUM_SEGMENTS, NUM_SEGMENTS, transform=tfms)
     print(f'AUC: {auc}')