delete LCAE, adjust training

cfg.py

@@ -2,8 +2,8 @@ from pathlib import Path
 import torch
 
 BATCH_SIZE = 128
-NUM_EPOCHS = 50
+NUM_EPOCHS = 1
-NUM_WORKERS = 4
+NUM_WORKERS = 0
 NUM_SEGMENTS = 5
 NUM_SEGMENT_HOPS = 2
 SEEDS = [42, 1337]

main.py

@@ -2,7 +2,7 @@ import numpy as np
 from tqdm import tqdm
 from cfg import *
 from mimii import MIMII
-from models.ae import AE, LCAE
+from models.ae import AE
 import torch.nn as nn
 import torch.optim as optim
 import random
@@ -12,10 +12,11 @@ torch.cuda.manual_seed(42)
 np.random.seed(42)
 random.seed(42)
 
-dataset_path = ALL_DATASET_PATHS[5]
+dataset_path = ALL_DATASET_PATHS[0]
 print(f'Training on {dataset_path.name}')
-mimii = MIMII(dataset_path=ALL_DATASET_PATHS[5], machine_id=0)
+mimii = MIMII(dataset_path=dataset_path, machine_id=0)
-mimii.preprocess(n_fft=1024, hop_length=512, n_mels=64, center=False, power=2.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,
@@ -26,7 +27,7 @@ dl = mimii.train_dataloader(
 )
 
 
-model = LCAE(320).to(DEVICE)
+model = AE(400).to(DEVICE)
 model.init_weights()
 criterion = nn.MSELoss()
 optimizer = optim.Adam(model.parameters(), lr=0.001)
@@ -38,24 +39,19 @@ 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)
+        y_hat, y = model(data)
-        loss = criterion(preds, data) + beta_1*entropy.mean() + beta_2*diversity
+        loss = criterion(y_hat, y)
 
         optimizer.zero_grad()
         loss.backward()
         optimizer.step()
-        #print(reconstruction.shape)
+
         losses.append(loss.item())
-        entropies.append(entropy.mean().item())
+    print(f'Loss: {np.mean(losses)}')
-        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}')

mimii.py

@@ -32,6 +32,12 @@ class MIMII(object):
         self.train_paths = train[len(test):]
         self.test_paths = normal_test + test
 
+        self.device = 'cpu'
+
+    def to(self, device):
+        self.device = device
+        return self
+
     def _normalize(self, S):
         return np.clip((S - self.min_level_db) / -self.min_level_db, 0, 1)
 
@@ -52,6 +58,7 @@ class MIMII(object):
                 mel_spec_norm = self._normalize(mel_spec_db)
                 m, n = mel_spec_norm.shape
                 np.save(folder/(file.stem + f'_{m}_{n}.npy'), mel_spec_norm)
+        return self
 
     def train_dataloader(self, segment_len=20, hop_len=5, **kwargs):
         # return both!!!
@@ -83,11 +90,10 @@ class MIMII(object):
                 file_preds = []
                 for batch in loader:
                     data, labels = batch
-                    data = data.to('cuda')
+                    data = data.to(self.device)
-                    data = data.view(data.shape[0], -1)
 
-                    y_hat, entropy, diversity = f(data)
+                    y_hat, y = f(data)
-                    preds = torch.sum((y_hat - data) ** 2, dim=tuple(range(1, y_hat.dim())))
+                    preds = torch.sum((y_hat - y) ** 2, dim=tuple(range(1, y_hat.dim())))
 
                     file_preds += preds.cpu().data.tolist()
                 y_true.append(labels.max().item())

models/ae.py

@@ -3,7 +3,7 @@ import torch.nn as nn
 import torch.functional as F
 
 class AE(nn.Module):
-    def __init__(self, in_dim=320):
+    def __init__(self, in_dim=400):
         super(AE, self).__init__()
         self.net = nn.Sequential(
             nn.Linear(in_dim, 64),
@@ -16,69 +16,13 @@ class AE(nn.Module):
             nn.ReLU(),
             nn.Linear(64, 64),
             nn.ReLU(),
-            nn.Linear(64, 320),
+            nn.Linear(64, in_dim),
             nn.ReLU(),
         )
 
     def forward(self, data):
-        return self.net(data)
+        x = data.view(data.shape[0], -1)
-
+        return self.net(x), x
-    def init_weights(self):
-        def _weight_init(m):
-            if hasattr(m, 'weight'):
-                if isinstance(m.weight, torch.Tensor):
-                    torch.nn.init.xavier_uniform_(m.weight,
-                                                  gain=nn.init.calculate_gain('relu'))
-            if hasattr(m, 'bias'):
-                if isinstance(m.bias, torch.Tensor):
-                    m.bias.data.fill_(0.01)
-
-        self.apply(_weight_init)
-
-
-
-class LCAE(nn.Module):
-    def __init__(self, in_dim=320):
-        super(LCAE, self).__init__()
-        num_mem = 10
-        mem_size= 8
-        self.num_mem = num_mem
-        self.encode = nn.Sequential(
-            nn.Linear(in_dim, 64),
-            nn.ReLU(),
-            nn.Linear(64, 64),
-            nn.ReLU(),
-            nn.Linear(64, num_mem),
-            nn.Softmax(-1)
-        )
-
-        self.decode = nn.Sequential(
-            nn.Linear(mem_size, 64),
-            nn.ReLU(),
-            nn.Linear(64, 64),
-            nn.ReLU(),
-            nn.Linear(64, 320),
-            nn.ReLU(),
-        )
-
-        self.M = nn.Parameter(
-            torch.randn(num_mem, mem_size)
-        )
-
-    def forward(self, data):
-        alphas = self.encode(data).unsqueeze(-1)
-        entropy_alphas = (alphas * -alphas.log()).sum(1)
-        M = self.M.expand(data.shape[0], *self.M.shape)
-        #print(M.shape, alphas.shape)  # torch.Size([128, 4, 8]) torch.Size([128, 4, 1])
-        elu = nn.ELU()
-        weighted = alphas * (1+elu(M+1e-13))
-        #print(weighted.shape)
-        summed = weighted.sum(1)
-        #print(summed.shape)
-        decoded = self.decode(summed)
-        diversity = (alphas.sum(dim=0)/data.shape[0]).max()
-        #print(alphas[0])
-        return decoded, entropy_alphas, diversity
 
     def init_weights(self):
         def _weight_init(m):

models/layers.py

@@ -21,8 +21,8 @@ class Subspectrogram(object):
 
 if __name__ == '__main__':
     import numpy as np
-    sub_spec_tnfm = Subspectrogram(20, 10)
+    sub_spec_tnfm = Subspectrogram(20, 20)
-    X = np.random.rand(1, 60, 40)
+    X = np.random.rand(1, 80, 40)
     Y = sub_spec_tnfm(X)
-    print(f'\t Sub-Spectrogram transformation from shape {X.shape} to {Y.shape}')
+    print(f'\tSub-Spectrogram transformation from shape {X.shape} to {Y.shape}')
-    print('Done ...')
+    print('\tDone ...')