delete LCAE, adjust training
This commit is contained in:
parent
3d7dbf222c
commit
55402d219c
4
cfg.py
4
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]
|
||||
|
24
main.py
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}')
|
||||
|
14
mimii.py
14
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())
|
||||
|
64
models/ae.py
64
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):
|
||||
|
@ -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 ...')
|
||||
print(f'\tSub-Spectrogram transformation from shape {X.shape} to {Y.shape}')
|
||||
print('\tDone ...')
|
Loading…
x
Reference in New Issue
Block a user