delete LCAE, adjust training

This commit is contained in:
Robert Müller
2020-03-18 17:13:00 +01:00
parent 3d7dbf222c
commit 55402d219c
5 changed files with 30 additions and 84 deletions

24
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}')