initial commit

main.py

@@ -0,0 +1,64 @@
+import numpy as np
+from tqdm import tqdm
+from cfg import *
+from mimii import MIMII
+from models.ae import AE, LCAE
+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)
+
+dataset_path = ALL_DATASET_PATHS[5]
+print(f'Training on {dataset_path.name}')
+mimii = MIMII(dataset_path=ALL_DATASET_PATHS[5], machine_id=0)
+mimii.preprocess(n_fft=1024, hop_length=512, n_mels=64, 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
+)
+
+
+model = LCAE(320).to(DEVICE)
+model.init_weights()
+criterion = nn.MSELoss()
+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 = []
+    entropies = []
+    l1s = []
+    for batch in tqdm(dl):
+        data, labels = batch
+        data = data.to(DEVICE)
+        data = data.view(data.shape[0], -1)
+
+        preds, entropy, diversity = model(data)
+        loss = criterion(preds, data) + beta_1*entropy.mean() + beta_2*diversity
+
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+        #print(reconstruction.shape)
+        losses.append(loss.item())
+        entropies.append(entropy.mean().item())
+        l1s.append(diversity.item())
+    print(f'Loss: {np.mean(losses)}; Entropy: {np.mean(entropies)}; l1:{np.mean(l1s)}')
+
+auc = mimii.evaluate_model(model, NUM_SEGMENTS, NUM_SEGMENTS)
+print(f'AUC: {auc}')
+
+
+