add more modularity

main.py

@@ -1,60 +1,57 @@
-import numpy as np
-from tqdm import tqdm
-from cfg import *
-from mimii import MIMII
-from models.ae import AE
-import torch.nn as nn
-import torch.optim as optim
-import random
+if __name__ == '__main__':
+    import numpy as np
+    from tqdm import tqdm
+    from cfg import *
+    from mimii import MIMII
+    from models.ae import AE
+    import torch.nn as nn
+    import torch.optim as optim
+    import random
 
-torch.manual_seed(42)
-torch.cuda.manual_seed(42)
-np.random.seed(42)
-random.seed(42)
+    torch.manual_seed(42)
+    torch.cuda.manual_seed(42)
+    np.random.seed(42)
+    random.seed(42)
 
-dataset_path = ALL_DATASET_PATHS[0]
-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)
+    dataset_path = ALL_DATASET_PATHS[0]
+    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)
 
-dl = mimii.train_dataloader(
-    segment_len=NUM_SEGMENTS,
-    hop_len=NUM_SEGMENT_HOPS,
-    batch_size=BATCH_SIZE,
-    num_workers=NUM_WORKERS,
-    shuffle=True
-)
+    dl = mimii.train_dataloader(
+        segment_len=NUM_SEGMENTS,
+        hop_len=NUM_SEGMENT_HOPS,
+        batch_size=BATCH_SIZE,
+        num_workers=NUM_WORKERS,
+        shuffle=True
+    )
 
 
-model = AE(400).to(DEVICE)
-model.init_weights()
-criterion = nn.MSELoss()
-optimizer = optim.Adam(model.parameters(), lr=0.001)
+    model = AE(400).to(DEVICE)
+    model.init_weights()
+
+    optimizer = optim.Adam(model.parameters(), lr=0.001)
 
 
-beta_1 = 0.00
-beta_2 = 0.0
+    for epoch in range(NUM_EPOCHS):
+        print(f'EPOCH #{epoch+1}')
+        losses = []
+        for batch in tqdm(dl):
+            data, labels = batch
+            data = data.to(DEVICE)
 
-for epoch in range(NUM_EPOCHS):
-    print(f'EPOCH #{epoch+1}')
-    losses = []
-    for batch in tqdm(dl):
-        data, labels = batch
-        data = data.to(DEVICE)
+            loss = model.train_loss(data)
 
-        y_hat, y = model(data)
-        loss = criterion(y_hat, y)
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
 
-        optimizer.zero_grad()
-        loss.backward()
-        optimizer.step()
+            losses.append(loss.item())
+        print(f'Loss: {np.mean(losses)}')
 
-        losses.append(loss.item())
-    print(f'Loss: {np.mean(losses)}')
-
-auc = mimii.evaluate_model(model, NUM_SEGMENTS, NUM_SEGMENTS)
-print(f'AUC: {auc}')
+    auc = mimii.evaluate_model(model, NUM_SEGMENTS, NUM_SEGMENTS)
+    print(f'AUC: {auc}')