working commit 2

cfg.py

@@ -1,15 +1,15 @@
 from pathlib import Path
 import torch
 
-BATCH_SIZE = 128
-NUM_EPOCHS = 50
+BATCH_SIZE = 96
+NUM_EPOCHS = 100
 NUM_WORKERS = 4
 NUM_SEGMENTS = 80
-NUM_SEGMENT_HOPS = 20
-SEEDS = [42, 1337]
+NUM_SEGMENT_HOPS = 15
+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,22 +1,24 @@
 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)
-    torch.cuda.manual_seed(42)
-    np.random.seed(42)
-    random.seed(42)
+    def train(dataset_path, machine_id, band, norm, seed):
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed(seed)
+        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}')
-    mimii = MIMII(dataset_path=dataset_path, machine_id=0)
+        mimii = MIMII(dataset_path=dataset_path, machine_id=machine_id)
         mimii.to(DEVICE)
         #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)
@@ -30,7 +32,7 @@ if __name__ == '__main__':
             transform=tfms
         )
 
-    model = SubSpecCAE().to(DEVICE)
+        model = SubSpecCAE(norm=norm, band=band).to(DEVICE)
         model.init_weights()
 
         # print(model(torch.randn(128, 1, 20, 80).to(DEVICE)).shape)
@@ -55,7 +57,24 @@ if __name__ == '__main__':
             print(f'Loss: {np.mean(losses)}')
 
         auc = mimii.evaluate_model(model, NUM_SEGMENTS, NUM_SEGMENTS, transform=tfms)
-    print(f'AUC: {auc}')
+        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)
+
+
+
+
 
 
 

mimii.py

@@ -11,9 +11,7 @@ __all__ = ['MIMII']
 
 
 class MIMII(object):
-    def __init__(self, dataset_path, machine_id, seed=42):
-        torch.random.manual_seed(seed)
-        np.random.seed(seed)
+    def __init__(self, dataset_path, machine_id):
         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):
-            sub_spec = sample[:, i:i+self.hop_size:,]
+        neg = 0 if self.height == self.hop_size else 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)