65 lines
1.6 KiB
Python
65 lines
1.6 KiB
Python
import numpy as np
|
|
from tqdm import tqdm
|
|
from cfg import *
|
|
from mimii import MIMII
|
|
from models.ae import AE, LCAE
|
|
import torch.nn as nn
|
|
import torch.optim as optim
|
|
import random
|
|
|
|
torch.manual_seed(42)
|
|
torch.cuda.manual_seed(42)
|
|
np.random.seed(42)
|
|
random.seed(42)
|
|
|
|
dataset_path = ALL_DATASET_PATHS[5]
|
|
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)
|
|
|
|
dl = mimii.train_dataloader(
|
|
segment_len=NUM_SEGMENTS,
|
|
hop_len=NUM_SEGMENT_HOPS,
|
|
batch_size=BATCH_SIZE,
|
|
num_workers=NUM_WORKERS,
|
|
shuffle=True
|
|
)
|
|
|
|
|
|
model = LCAE(320).to(DEVICE)
|
|
model.init_weights()
|
|
criterion = nn.MSELoss()
|
|
optimizer = optim.Adam(model.parameters(), lr=0.001)
|
|
|
|
|
|
beta_1 = 0.00
|
|
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
|
|
|
|
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)}')
|
|
|
|
auc = mimii.evaluate_model(model, NUM_SEGMENTS, NUM_SEGMENTS)
|
|
print(f'AUC: {auc}')
|
|
|
|
|
|
|