diff --git a/cfg.py b/cfg.py index 8bc9e7d..a1cf078 100644 --- a/cfg.py +++ b/cfg.py @@ -2,8 +2,8 @@ from pathlib import Path import torch BATCH_SIZE = 128 -NUM_EPOCHS = 50 -NUM_WORKERS = 4 +NUM_EPOCHS = 1 +NUM_WORKERS = 0 NUM_SEGMENTS = 5 NUM_SEGMENT_HOPS = 2 SEEDS = [42, 1337] diff --git a/main.py b/main.py index cc82f6c..68a0543 100644 --- a/main.py +++ b/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.preprocess(n_fft=1024, hop_length=512, n_mels=64, center=False, power=2.0) +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, @@ -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) - loss = criterion(preds, data) + beta_1*entropy.mean() + beta_2*diversity + y_hat, y = model(data) + loss = criterion(y_hat, y) 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)}') + print(f'Loss: {np.mean(losses)}') auc = mimii.evaluate_model(model, NUM_SEGMENTS, NUM_SEGMENTS) print(f'AUC: {auc}') diff --git a/mimii.py b/mimii.py index 3e4d702..a5b8c0e 100644 --- a/mimii.py +++ b/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.view(data.shape[0], -1) + data = data.to(self.device) - y_hat, entropy, diversity = f(data) - preds = torch.sum((y_hat - data) ** 2, dim=tuple(range(1, y_hat.dim()))) + y_hat, y = f(data) + 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()) diff --git a/models/ae.py b/models/ae.py index 8b0f461..9c0dafc 100644 --- a/models/ae.py +++ b/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) - - 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 + x = data.view(data.shape[0], -1) + return self.net(x), x def init_weights(self): def _weight_init(m): diff --git a/models/layers.py b/models/layers.py index a346bc1..07f6f08 100644 --- a/models/layers.py +++ b/models/layers.py @@ -21,8 +21,8 @@ class Subspectrogram(object): if __name__ == '__main__': import numpy as np - sub_spec_tnfm = Subspectrogram(20, 10) - X = np.random.rand(1, 60, 40) + sub_spec_tnfm = Subspectrogram(20, 20) + 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('Done ...') \ No newline at end of file + print(f'\tSub-Spectrogram transformation from shape {X.shape} to {Y.shape}') + print('\tDone ...') \ No newline at end of file