steffen/audio_anomaly_detection
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

working commit 2

Browse Source
This commit is contained in:
Robert Müller 2020-03-20 09:27:49 +01:00
parent cc9e9b50a4
commit 0f325676e5
5 changed files with 69 additions and 50 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
cfg.py
View File

@ -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

95
main.py
View File

@ -1,61 +1,80 @@
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.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)
print(f'Training on {dataset_path.name}')
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)
dl = mimii.train_dataloader(
segment_len=NUM_SEGMENTS,
hop_len=NUM_SEGMENT_HOPS,
batch_size=BATCH_SIZE,
num_workers=NUM_WORKERS,
shuffle=True,
transform=tfms
)
dl = mimii.train_dataloader(
segment_len=NUM_SEGMENTS,
hop_len=NUM_SEGMENT_HOPS,
batch_size=BATCH_SIZE,
num_workers=NUM_WORKERS,
shuffle=True,
transform=tfms
)
model = SubSpecCAE().to(DEVICE)
model.init_weights()
model = SubSpecCAE(norm=norm, band=band).to(DEVICE)
model.init_weights()
# print(model(torch.randn(128, 1, 20, 80).to(DEVICE)).shape)
# print(model(torch.randn(128, 1, 20, 80).to(DEVICE)).shape)
optimizer = optim.Adam(model.parameters(), lr=0.001)
optimizer = optim.Adam(model.parameters(), lr=0.001)
for epoch in range(NUM_EPOCHS):
print(f'EPOCH #{epoch+1}')
losses = []
for batch in tqdm(dl):
data, labels = batch
data = data.to(DEVICE) # torch.Size([128, 4, 20, 80]) batch x subs_specs x height x width
for epoch in range(NUM_EPOCHS):
print(f'EPOCH #{epoch+1}')
losses = []
for batch in tqdm(dl):
data, labels = batch
data = data.to(DEVICE) # torch.Size([128, 4, 20, 80]) batch x subs_specs x height x width
loss = model.train_loss(data)
optimizer.zero_grad()
loss.backward()
optimizer.step()
losses.append(loss.item())
print(f'Loss: {np.mean(losses)}')
auc = mimii.evaluate_model(model, NUM_SEGMENTS, NUM_SEGMENTS, transform=tfms)
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)
loss = model.train_loss(data)
optimizer.zero_grad()
loss.backward()
optimizer.step()
losses.append(loss.item())
print(f'Loss: {np.mean(losses)}')
auc = mimii.evaluate_model(model, NUM_SEGMENTS, NUM_SEGMENTS, transform=tfms)
print(f'AUC: {auc}')

4
mimii.py
View File

@ -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}'

5
models/ae.py
View File

@ -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

5
models/layers.py
View File

@ -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)