【超初心者向け】AE(AutoEncoder)をPython(PyTorch)で実装してみる。

import os
import sys
import numpy as np
import librosa
import glob

##############################################################
'''
ここから別記事（https://tips-memo.com/python-mel）で解説済み
'''
##############################################################
def extract_mel(wav, sr, n_mels=64): #(timeframe, mel_dim)
    audio, _ = librosa.load(wav, sr=sr)
    mel = librosa.feature.melspectrogram(y=audio, sr=sr, n_mels=n_mels).T
    return mel

def data_list(path, sr, n_mels=64):
    wav_list = glob.glob(path)
    size = len(wav_list)
    data = np.ones((1, n_mels))
    count= 0
    for wavname in wav_list:
        component = extract_mel(wavname, sr=sr, n_mels=64)
        data = np.concatenate([data, component], axis=0)
        count += 1
        sys.stdout.write("\r%s" % "現在"+str(np.around((count/len(wav_list))*100 , 2))+"%完了")
        sys.stdout.flush()
    return data[1:], size
##############################################################
'''
ここまで別記事（https://tips-memo.com/python-mel）で解説済み
'''
##############################################################

Toytype = "ToyCar"

name = "exp1_dataset_" + Toytype

# pathの設定
train_wav_path= "./ToyADMOS-dataset-master/{}/train_normal/*".format(name)
test_normal_wav_path = "./ToyADMOS-dataset-master/{}/test_normal/*".format(name)
test_anomaly_wav_path = "./ToyADMOS-dataset-master/{}/test_anomaly/*".format(name)

# パラメータの保存先指定
out_audio_dir = "../data/audio/ToyADMOS/{}".format(Toytype)
if not os.path.exists(out_audio_dir):
    os.makedirs(out_audio_dir)

# wavデータの一括読み込み
# 今回はメル周波数スペクトログラムを16000Hz, 64次元取得して全データに関して縦に並べている
print("---read_training_wav---")
train_data_list, size_train = my_modules.data_list(train_wav_path, sr=16000, n_fft=512, hop_length=160, n_mels=64)
print("\n---read_test_normal_wav---")
test_normal_data_list, size_test_normal = my_modules.data_list(test_normal_wav_path, sr=16000, n_fft=512, hop_length=160, n_mels=64)
print("\n---read_test_anomaly_wav---")
test_anomaly_data_list, size_test_anomaly = my_modules.data_list(test_anomaly_wav_path, sr=16000, n_fft=512, hop_length=160, n_mels=64)
print("\n")

# メル周波数スペクトログラムの保存
np.save('./{}/mel/train_mel_16000Hz_64dim.npy'.format(out_audio_dir), train_data_list)
np.save('./{}/mel/test_normal_mel_16000Hz_64dim.npy'.format(out_audio_dir), test_normal_data_list)
np.save('./{}/mel/test_anomaly_mel_16000Hz_64dim.npy'.format(out_audio_dir), test_anomaly_data_list)

# サイズの保存
np.save('./{}/size/size_train.npy'.format(out_audio_dir), np.array([size_train]))
np.save('./{}/size/size_test_normal.npy'.format(out_audio_dir), np.array([size_test_normal]))
np.save('./{}/size/size_test_anomaly.npy'.format(out_audio_dir), np.array([size_test_anomaly]))

import os
import numpy as np
import torch
import torchvision
from torch import nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import transforms
from torch.optim.lr_scheduler import LambdaLR
import matplotlib.pyplot as plt

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model_name = "AE"
dataset_type = "ToyADMOS"
Toy_type = "ToyCar" #"ToyCar" ot "ToyConveyor" or "ToyTrain"

out_params_dir = "../data/models/{}/params/{}/{}".format(model_name, dataset_type, Toy_type)
out_figs_dir = "../data/models/{}/images/{}/{}".format(model_name, dataset_type,Toy_type)
out_texts_dir = "../data/models/{}/texts/{}/{}".format(model_name,dataset_type,Toy_type)
in_audio_dir = "../data/audio/{}/{}".format(dataset_type,Toy_type)

if not os.path.exists(out_params_dir):
    os.makedirs(out_params_dir)
if not os.path.exists(out_figs_dir):
    os.makedirs(out_figs_dir)
if not os.path.exists(out_texts_dir):
    os.makedirs(out_texts_dir)
if not os.path.exists(in_audio_dir):
    os.makedirs(in_audio_dir)

# データのロード
train_data_list = np.load('./{}/mel/train_mel_16000Hz_64dim.npy'.format(in_audio_dir))

# データのサイズを取得・表示
size_train = np.load('./{}/size/size_train.npy'.format(in_audio_dir))[0]
print("\n訓練用データサイズ："+str(size_train))

from torch.utils.data import Dataset

# ExpandDatasetは1つの時間フレームに対して前後10フレームを取得するようなデータセット
class ExpandDataset(Dataset):
    def __init__(self, data, transform=None):
        self.transform = transform
        self.data = data
        self.data_num = len(data)
        if self.data.ndim==2:
            self.pad_data_fr = data[:10][::-1]
            self.pad_data_bc = data[-10:][::-1]
            self.pad_data = np.concatenate([self.pad_data_fr, data, self.pad_data_bc], axis=0)
        elif self.data.ndim==3:
            self.pad_data_fr = data[:,:10,:][:,::-1,:]
            self.pad_data_bc = data[:,-10:,:][:,::-1,:]
            self.pad_data = np.concatenate([self.pad_data_fr, data, self.pad_data_bc], axis=1)    
    def __len__(self):
        return self.data_num
    def __getitem__(self, idx):
        if self.transform:
            if self.data.ndim==2:
                out_data = self.transform(self.pad_data)[0][idx:idx+20].flatten()
            elif self.data.ndim==3:
                index = int(random.uniform(0,self.data.shape[1]))
                out_data = self.transform(self.pad_data)[:,idx, index:index+20].flatten()
        else:
            print("transformを使用しテンソル化してください")
        return out_data

class Autoencoder(nn.Module):
    def __init__(self):
        super(Autoencoder, self).__init__()
        self.dense_enc1 = nn.Linear(1280, 512)
        self.bn1 = nn.BatchNorm1d(512)
        self.dense_enc2 = nn.Linear(512, 256)
        self.bn2 = nn.BatchNorm1d(256)
        self.dense_enc3 = nn.Linear(256,128)
    
        self.dense_dec1 = nn.Linear(128,256)
        self.bn4 = nn.BatchNorm1d(256)
        self.dense_dec2 = nn.Linear(256, 512)
        self.bn5 = nn.BatchNorm1d(512)
        self.drop1 = nn.Dropout(p=0.2)
        self.dense_dec3 = nn.Linear(512, 1280)

    def encoder(self, x):
        x = F.relu(self.dense_enc1(x))
        x = self.bn1(x)
        x = F.relu(self.dense_enc2(x))
        x = self.bn2(x)
        x = self.dense_enc3(x)
        return x

    def decoder(self, x):
        x = F.relu(self.dense_dec1(x))
        x = self.bn4(x)
        x = F.relu(self.dense_dec2(x))
        x = self.bn5(x)
        x = self.drop1(x)
        x = self.dense_dec3(x)
        return x

    def forward(self, x):
        z = self.encoder(x)
        x = self.decoder(z)
        return x, z

transform = transforms.Compose([transforms.ToTensor()])
train_dataset = my_dataset.ExpandDataset(train_data_list, transform)

train_batch_size = 100
train_loader = DataLoader(train_dataset, batch_size=train_batch_size, shuffle=True)

num_epochs = 200
learning_rate = 1e-4

model = Autoencoder().to(device)
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=1e-4)

def func_100(epoch):
    if epoch <= 100:
        return 1
    elif 100 < epoch:
        return -0.99*(1e-2)*(epoch) + 1.99

scheduler = LambdaLR(optimizer, lr_lambda=my_scheduler.func_100)

model.train()
loss_list = []
for epoch in range(num_epochs):
    losses = []
    for x in train_loader:
        x = x.to(device)
        model.zero_grad()
        y, z = model(x.float())
        loss = criterion(y.float(), x.float())
        loss.backward()
        optimizer.step()
        losses.append(loss.cpu().detach().numpy())
    print('epoch [{}/{}], loss: {:.4f}'.format(epoch + 1, num_epochs, np.average(losses))) 
    scheduler.step()
    loss_list.append(np.average(losses))

name_loss = "train_loss_"+str(model_name)+"_"+str(feature_type)+".npy"
name_params = "params_"+str(model_name)+"_"+str(feature_type)+".pth"
name_graph = "train_loss_"+str(model_name)+"_"+str(feature_type)+".png"
np.save("./{}/{}".format(out_params_dir, name_loss), np.array(loss_list))
torch.save(model.state_dict(), './{}/{}'.format(out_params_dir, name_params))

plt.figure()
plt.plot(range(epoch), loss_list[1:], 'r-', label='train_loss')
plt.legend()
plt.xlabel('epoch')
plt.ylabel('loss')
plt.grid()
plt.savefig('./{}/{}'.format(out_figs_dir, name_graph))
plt.close()

# データの保存先と読み込み先の指定
name_roc = "roc_curve_"+str(model_name)+"_"+str(feature_type)+".png"
name_y_true = "y_true_" +str(model_name)+"_"+str(feature_type)
name_y_score = "y_score_" +str(model_name)+"_"+str(feature_type)

# テストデータの読み込み
test_normal_data_list =　np.load('./{}/mel/test_normal_mel_16000Hz_64dim.npy'.format(in_audio_dir))
test_anomaly_data_list = np.load('./{}/mel/test_anomaly_mel_16000Hz_64dim.npy'.format(in_audio_dir))

# データサイズの出力
size_test_normal = np.load('./{}/size/size_test_normal.npy'.format(in_audio_dir))[0]
size_test_anomaly = np.load('./{}/size/size_test_anomaly.npy'.format(in_audio_dir))[0]
print("\n正常音テストデータサイズ："+str(size_test_normal))
print("異常音テストデータサイズ："+str(size_test_anomaly))

# バッチサイズの指定
# ここではデータサイズの約数になるように定めています
test_normal_batch_size = 70
test_anomaly_batch_size = 66

# データセットの準備
transform = transforms.Compose([transforms.ToTensor()])
test_normal_dataset = my_dataset.ExpandDataset(test_normal_data_list, transform)
test_anomaly_dataset = my_dataset.ExpandDataset(test_anomaly_data_list, transform)
test_normal_loader = DataLoader(test_normal_dataset, batch_size=test_normal_batch_size, shuffle=False)
test_anomaly_loader = DataLoader(test_anomaly_dataset, batch_size=test_anomaly_batch_size, shuffle=False)

# テスト開始
model = model.to(device)
params = torch.load('./{}/{}'.format(out_params_dir, name_params))
model.load_state_dict(params)
model.eval()

# lossを1フレームごとに算出するためにcriterionを再設定
criterion = nn.MSELoss(reduction="none")

# 正常データテスト
losses_normal = []
count=0
for x in test_normal_loader:
    x = x.to(device)
    y, z = model(x.float())
    loss_normal = criterion(y.float(), x.float())
    losses_normal.append(loss_normal.cpu().detach().numpy())
    sys.stdout.write("\r%s" % "現在"+str(np.around((count/len(test_normal_loader))*100 , 2))+"%完了")
    sys.stdout.flush()
    count += 1

# 異常スコアの計算
# まずは1280の方向に平均を取って一次元化する
output_normal = np.mean(np.array(losses_normal), axis=2).flatten()
# wavファイルごとに分割する
output_normal_split = np.array(np.split(output_normal, size_test_normal))
# 最初と最後の50フレームを削除する
# 理由はINDファイルを利用しており，最初と最後には合成した背景雑音しか含まれていないため
output_normal_split_cut = output_normal_split[:,50:-50]
# 各wavファイルごとに異常度の最大値をスコアとして採用する
output_normal_max = output_normal_split_cut.max(axis=1)

# 異常データテスト
losses_anomaly = []
count=0
for x in test_anomaly_loader:
    x = x.to(device)
    y, z = model(x.float())
    loss_anomaly = criterion(y.float(), x.float())
    losses_anomaly.append(loss_anomaly.cpu().detach().numpy())
    sys.stdout.write("\r%s" % "現在"+str(np.around((count/len(test_anomaly_loader))*100 , 2))+"%完了")
    sys.stdout.flush()
    count += 1

# 異常スコアの計算
# まずは1280の方向に平均を取って一次元化する
output_anomaly = np.mean(np.array(losses_anomaly), axis=2).flatten()
# wavファイルごとに分割する
output_anomaly_split = np.array(np.split(output_anomaly, size_test_anomaly))
# 最初と最後の50フレームを削除する
# 理由はINDファイルを利用しており，最初と最後には合成した背景雑音しか含まれていないため
output_anomaly_split_cut = output_anomaly_split[:,50:-50]
# 各wavファイルごとに異常度の最大値をスコアとして採用する
output_anomaly_max = output_anomaly_split_cut.max(axis=1)

# ROC曲線の描画
y_true = np.concatenate([np.zeros(size_test_normal), np.ones(size_test_anomaly)])
y_score = np.concatenate([output_normal_max, output_anomaly_max])
np.save("./{}/{}".format(out_params_dir, name_y_true), y_true)
np.save("./{}/{}".format(out_params_dir, name_y_score), y_score)
fpr, tpr, thresholds = metrics.roc_curve(y_true, y_score)
auc = metrics.auc(fpr, tpr)
plt.plot(fpr, tpr, label='ROC curve (area = %.2f)'%auc)
plt.legend()
plt.xlabel('FPR: False positive rate')
plt.ylabel('TPR: True positive rate')
plt.grid()
plt.savefig("./{}/{}".format(out_figs_dir, name_roc))
plt.close()

# F値の計算
# FPRを0.1に抑えた時の閾値を採用する
upper_fpr = 0.1
fpr_index = np.argmin(np.abs(fpr - upper_fpr))
gap = y_score - thresholds[fpr_index]
y_pred = (gap>0)*1
f_score = f1_score(y_true, y_pred) 
thres = thresholds[fpr_index]
print("F-score under fpr=0.1:"+str(f_score))
print("thresholds under fpr=0.1:"+str(thres))

Out:
F-score_under_fpr=0.1：0.7927565392354124
Threshold：2.8711948