working commit 2

cfg.py

@@ -1,15 +1,15 @@
 from pathlib import Path
 import torch
 
-BATCH_SIZE = 128
+BATCH_SIZE = 96
-NUM_EPOCHS = 50
+NUM_EPOCHS = 100
 NUM_WORKERS = 4
 NUM_SEGMENTS = 80
-NUM_SEGMENT_HOPS = 20
+NUM_SEGMENT_HOPS = 15
-SEEDS = [42, 1337]
+SEEDS = [42, 1337, 666, 2012, 1e42]
 ALL_DATASET_PATHS = list((Path(__file__).parent.absolute() / 'data' / 'mimii').glob('*/'))
 
 DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 SUB_SPEC_HEIGT = 20
-SUB_SPEC_HOP = SUB_SPEC_HEIGT
+SUB_SPEC_HOP = 10

main.py

@@ -1,61 +1,80 @@
 if __name__ == '__main__':
     import numpy as np
+    import random
     from tqdm import tqdm
     from cfg import *
     from mimii import MIMII
     from models.ae import AE, SubSpecCAE
-    import torch.nn as nn
+    import pickle
     import torch.optim as optim
-    import random
     from models.layers import Subspectrogram
 
-    torch.manual_seed(42)
+    def train(dataset_path, machine_id, band, norm, seed):
-    torch.cuda.manual_seed(42)
+        torch.manual_seed(seed)
-    np.random.seed(42)
+        torch.cuda.manual_seed(seed)
-    random.seed(42)
+        np.random.seed(seed)
+        torch.cuda.manual_seed_all(seed)
+        torch.backends.cudnn.deterministic = True
+        random.seed(seed)
 
-    dataset_path = ALL_DATASET_PATHS[5]
+        print(f'Training on {dataset_path.name}')
-    print(f'Training on {dataset_path.name}')
+        mimii = MIMII(dataset_path=dataset_path, machine_id=machine_id)
-    mimii = MIMII(dataset_path=dataset_path, machine_id=0)
+        mimii.to(DEVICE)
-    mimii.to(DEVICE)
+        #mimii.preprocess(n_fft=1024, hop_length=256, n_mels=80, center=False, power=2.0)  # 80 x 80
-    #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)
-    tfms = Subspectrogram(SUB_SPEC_HEIGT, SUB_SPEC_HOP)
 
-    dl = mimii.train_dataloader(
+        dl = mimii.train_dataloader(
-        segment_len=NUM_SEGMENTS,
+            segment_len=NUM_SEGMENTS,
-        hop_len=NUM_SEGMENT_HOPS,
+            hop_len=NUM_SEGMENT_HOPS,
-        batch_size=BATCH_SIZE,
+            batch_size=BATCH_SIZE,
-        num_workers=NUM_WORKERS,
+            num_workers=NUM_WORKERS,
-        shuffle=True,
+            shuffle=True,
-        transform=tfms
+            transform=tfms
-    )
+        )
 
-    model = SubSpecCAE().to(DEVICE)
+        model = SubSpecCAE(norm=norm, band=band).to(DEVICE)
-    model.init_weights()
+        model.init_weights()
 
-    # print(model(torch.randn(128, 1, 20, 80).to(DEVICE)).shape)
+        # print(model(torch.randn(128, 1, 20, 80).to(DEVICE)).shape)
 
-    optimizer = optim.Adam(model.parameters(), lr=0.001)
+        optimizer = optim.Adam(model.parameters(), lr=0.001)
 
 
-    for epoch in range(NUM_EPOCHS):
+        for epoch in range(NUM_EPOCHS):
-        print(f'EPOCH #{epoch+1}')
+            print(f'EPOCH #{epoch+1}')
-        losses = []
+            losses = []
-        for batch in tqdm(dl):
+            for batch in tqdm(dl):
-            data, labels = batch
+                data, labels = batch
-            data = data.to(DEVICE)  # torch.Size([128, 4, 20, 80]) batch x subs_specs x height x width
+                data = data.to(DEVICE)  # torch.Size([128, 4, 20, 80]) batch x subs_specs x height x width
+
+                loss = model.train_loss(data)
+
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+
+                losses.append(loss.item())
+            print(f'Loss: {np.mean(losses)}')
+
+        auc = mimii.evaluate_model(model, NUM_SEGMENTS, NUM_SEGMENTS, transform=tfms)
+        print(f'AUC: {auc}, Machine: {machine_id}, Band: {band}, Norm: {norm}, Seed: {seed}')
+        return auc
+
+
+    results = []
+    for norm in ('instance', 'batch'):
+        for seed in SEEDS:
+            for dataset_path in ALL_DATASET_PATHS:
+                for machine_id in [0, 2, 4, 6]:
+                    for band in range(7):
+                        auc = train(dataset_path, machine_id, band, norm, seed)
+                        results.append([dataset_path.name, machine_id, seed, band, norm, auc])
+                        with open(f'results_{norm}.pkl', 'wb') as f:
+                            pickle.dump(results, f)
 
-            loss = model.train_loss(data)
 
-            optimizer.zero_grad()
-            loss.backward()
-            optimizer.step()
 
-            losses.append(loss.item())
-        print(f'Loss: {np.mean(losses)}')
 
-    auc = mimii.evaluate_model(model, NUM_SEGMENTS, NUM_SEGMENTS, transform=tfms)
-    print(f'AUC: {auc}')
 
 
 

mimii.py

@@ -11,9 +11,7 @@ __all__ = ['MIMII']
 
 
 class MIMII(object):
-    def __init__(self, dataset_path, machine_id, seed=42):
+    def __init__(self, dataset_path, machine_id):
-        torch.random.manual_seed(seed)
-        np.random.seed(seed)
         self.machine = dataset_path.name
         self.machine_id = machine_id
         self.root = dataset_path / f'id_0{machine_id}'

models/ae.py

@@ -67,11 +67,12 @@ class AE(nn.Module):
 
 
 class SubSpecCAE(nn.Module):
-    def __init__(self, F=20, T=80, norm='batch', activation='relu', dropout_prob=0.25):
+    def __init__(self, F=20, T=80, norm='batch', activation='relu', dropout_prob=0.25, band=0):
         super(SubSpecCAE, self).__init__()
         self.T = T
         self.F = F
         self.activation = activation
+        self.band = band
         Norm = nn.BatchNorm2d if norm == 'batch' else nn.InstanceNorm2d
         Activation = nn.ReLU if activation == 'relu' else nn.LeakyReLU
         self.encoder = nn.Sequential(
@@ -101,7 +102,7 @@ class SubSpecCAE(nn.Module):
         )
 
     def forward(self, x):
-        x = x[:,3,:,].unsqueeze(1)  # select a single supspec
+        x = x[:,self.band,:,].unsqueeze(1)  # select a single supspec
         encoded = self.encoder(x)
         decoded = self.decoder(encoded)
         return decoded, x

models/layers.py

@@ -13,8 +13,9 @@ class Subspectrogram(object):
             sample = sample.reshape(1, *sample.shape)
         # sample shape: 1 x num_mels x num_frames
         sub_specs = []
-        for i in range(0, sample.shape[1], self.hop_size):
+        neg = 0 if self.height == self.hop_size else self.hop_size
-            sub_spec = sample[:, i:i+self.hop_size:,]
+        for i in range(0, sample.shape[1]-neg, self.hop_size):
+            sub_spec = sample[:, i:i+self.height:,]
             sub_specs.append(sub_spec)
         return np.concatenate(sub_specs)