generator_eval.py

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import math
import sys
import librosa
from random import shuffle
import math
import time
from numpy import genfromtxt
import torch
from torch import nn
import torch.nn.functional as F
import torch.optim as optim
from torch.optim.lr_scheduler import StepLR
from torch.autograd import Variable
from torchsummary import summary
import soundfile as sf

import os, glob
pd.set_option('display.max_rows', 500)
import h5py
import pickle
from sklearn import preprocessing
import argparse
import logging
from sklearn.preprocessing import label_binarize
from statistics import mean, variance, median
from collections import Counter
import config

torch.backends.cudnn.enabled=False

sys.path.insert(1, os.path.join(sys.path[0], './utils'))
from utilities import (read_audio, create_folder,
                       get_filename, create_logging, calculate_accuracy,
                       print_accuracy, calculate_confusion_matrix,
                       move_data_to_gpu, audio_unify,
                       CWLoss_iemocap, get_model_iemocap, normalize_function,
                       CrossEntropyLoss, get_lr, set_lr, set_cyclic_lr)

# generator
sys.path.insert(1, os.path.join(sys.path[0], './audio_models/waveunet'))
# from seanet import GeneratorSEANt
# from gnet import AudioGeneratorResnet
from waveunet import Waveunet
# audio pre-processing
from transformers import Wav2Vec2FeatureExtractor, AutoConfig, PretrainedConfig
feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained("facebook/wav2vec2-base-960h")
# hugging models
from transformers import AutoModelForAudioClassification, HubertForSequenceClassification, WavLMForSequenceClassification
# end2you audio models
from audio_rnn_model import AudioRNNModel


# For pytorch dataset
from torch.utils.data import TensorDataset, DataLoader
from datasets.dataset_dict import DatasetDict
from datasets import Dataset, load_metric

metric = load_metric("recall")
batch_size = config.batch_size
class_num = config.iemocap_num_classes
audio_len = config.iemocap_audio_samples
models_menu = config.iemocap_models_menu
criterion = CWLoss_iemocap

def data_generater(hdf5_path):
    '''Read data into a dict'''
    with h5py.File(hdf5_path, 'r') as hf:
        x_org = hf['x_org'][:]
        x_adv = hf['x_adv'][:]
        y = hf['y'][:]

    hf.close()

    # d = Dataset.from_dict({'label': y, 'x_org': x_org, 'x_adv': x_adv})
    d = {'y': y, 'x_org': x_org, 'x_adv': x_adv}
    return d

def forward_evaluate(model_t, data_loader, source_model_name, target_model_name, eps, cuda, z):
    # z is True if model_t is not zhao19 or emo18
    adv_acc = 0
    clean_acc = 0
    fool_rate = 0
    target_rate = 0
    norm = 0
    distortion = 0
    time_count = 0
    save_five = 0

    test_size = len(data_loader.dataset)
    assert data_loader.batch_size == 1

    for (idx, (batch_x_org, batch_x_adv, batch_y)) in enumerate(data_loader, 0):

        batch_x_org = move_data_to_gpu(batch_x_org, cuda)
        batch_x_adv = move_data_to_gpu(batch_x_adv, cuda)
        batch_y = move_data_to_gpu(batch_y, cuda)

        model_t.eval()
        if z:
            clean_out = model_t(normalize_function(batch_x_org.clone().detach()))[0]
        else:
            clean_out = model_t(normalize_function(batch_x_org.clone().detach()))

        clean_acc += torch.sum(clean_out.argmax(dim=-1) == batch_y).item()


        if z:
            adv_out = model_t(normalize_function(batch_x_adv.clone().detach()))[0]
        else:
            adv_out = model_t(normalize_function(batch_x_adv.clone().detach()))

        adv_acc +=torch.sum(adv_out.argmax(dim=-1) == batch_y).item()

        fool_rate += torch.sum(adv_out.argmax(dim=-1) != clean_out.argmax(dim=-1)).item()

        if args.target != -1:
            target = torch.LongTensor(batch_y.size(0))
            target.fill_(args.target)
            target = target.cuda()
            target_rate += torch.sum(adv_out.argmax(dim=-1) == target).item()

        # only take the success ones into SNR and norm; only save successful adversarial audios
        if torch.sum(adv_out.argmax(dim=-1) != clean_out.argmax(dim=-1)).item() > 0:
            # per = torch.squeeze(adv_inf.clone().detach() * adv_0.clone().detach(), dim=1)
            per = torch.squeeze(batch_x_adv.clone().detach() - batch_x_org.clone().detach())
            i_dis = torch.log10(torch.max(torch.abs(batch_x_org), dim=-1).values / torch.max(torch.abs(per), dim=-1).values)
            i_dis_no_inf = torch.where(i_dis == float('inf'), torch.tensor(0.0).cuda().detach(), i_dis.detach())
            distortion += torch.sum(20 * i_dis_no_inf).item()
            norm += torch.norm(per, 0)

            # Save some audio samples for human evaluation
            if save_five < 10 and test_size == 1151:
                cur_dis = torch.sum(20 * i_dis_no_inf).item() / 1
                cur_norm = torch.norm(per.detach(), 0) / 1
                audio_dir = os.path.join('/home/ychang/sparse_attack/workspace/audios_listen/IEMOCAP/generator2', '{}_{}'.format(source_model_name, target_model_name))
                create_folder(audio_dir)
                audio_org = os.path.join(audio_dir, '{:.4f}_{:.4f}_org.wav'.format(cur_dis, cur_norm))
                audio_adv = os.path.join(audio_dir, '{:.4f}_{:.4f}_adv.wav'.format(cur_dis, cur_norm))
                audio_per = os.path.join(audio_dir, '{:.4f}_{:.4f}_per.wav'.format(cur_dis, cur_norm))

                for i in range(1):
                    if batch_y[i] != torch.argmax(adv_out[i], dim=-1):
                        sf.write(audio_org, batch_x_org[i].clone().detach().cpu().numpy(), 16000)
                        sf.write(audio_adv, batch_x_adv[i].clone().detach().cpu().numpy(), 16000)
                        # per = np.squeeze(adv_0[i].clone().detach().cpu().numpy() * adv_inf[i].clone().detach().cpu().numpy(),0)
                        per = np.squeeze(batch_x_adv[i].clone().detach().cpu().numpy() - batch_x_org[i].clone().detach().cpu().numpy())
                        sf.write(audio_per, per, 16000)
                        save_five += 1
                        break

    if args.target != -1:
        logging.info('Clean: {0:.3%} Adversarial: {1:.3%} Fooling Rate: {2:.3%} Target Success Rate:{3:.3%}'.format(clean_acc/test_size, adv_acc/test_size, fool_rate/test_size, target_rate/test_size))
    else:
        # logging.info('Clean: {0:.3%} Adversarial: {1:.3%} Fooling Rate:{2:.3%}'.format(clean_acc/test_size, adv_acc/test_size, fool_rate/test_size))
        logging.info('fool rate: {:.3%}'.format(fool_rate/test_size))


def train(args):
    # Arugments & parameters
    workspace = args.workspace
    validation = args.validation
    epoch = args.epoch
    cuda = args.cuda
    source_model_name = args.source_model_name
    target_model_name = args.target_model_name
    rnn_name = args.rnn_name
    target = args.target
    eps = args.eps

    # adv samples preparation
    hdf5_adv = os.path.join(workspace, 'sparse_attack', 'sparse_adv', 'generator', 'iemocap_source_{}_val_{}.h5'.format(source_model_name, validation))
    data = data_generater(hdf5_adv)
    dataset = DatasetDict(data)
    # dataset = dataset.map(preprocess_function, remove_columns=["audio"], batched=True, batch_size=batch_size)
    dataset_eval = TensorDataset(torch.Tensor(dataset['x_org']), torch.Tensor(dataset['x_adv']), torch.LongTensor(dataset['y']))
    dataloader_eval = DataLoader(dataset_eval, batch_size=1, shuffle=False, num_workers=2)

    del data
    del dataset
    logging.info('{} adv samples loaded from {}'.format(len(dataloader_eval.dataset), hdf5_adv))

    #  models loading
    target_model, target_model_path = get_model_iemocap(target_model_name, validation, workspace)
    if 'finetune' not in target_model_path:
        target_model = AudioRNNModel(input_size=audio_len, num_outs=class_num, model_name=target_model_name, rnn_name=rnn_name)

    target_model.load_state_dict(torch.load(target_model_path))

    logging.info('the target eval model is {} and located in {}'.format(target_model_name, target_model_path))

    if cuda:
        target_model.cuda()
    target_model.eval()

    forward_evaluate(target_model, dataloader_eval, source_model_name, target_model_name, eps, cuda, target_model_name not in ['emo18', 'zhao19'])

    logging.info('finished evaluation for {}'.format(target_model_name))

if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='Example of parser. ')
    subparsers = parser.add_subparsers(dest='mode')

    parser_eval = subparsers.add_parser('eval')
    parser_eval.add_argument('--workspace', type=str, default='/storage/home/ychang/IEMOCAP')
    parser_eval.add_argument('--validation', action='store_true', default=False)
    parser_eval.add_argument('--epoch', type=int, required=False, default=20)
    parser_eval.add_argument('--cuda', action='store_true', default=False)
    parser_eval.add_argument('--eps', type=float, default=0.01, help='perturbation budget')
    parser_eval.add_argument('--source_model_name', type=str, choices=['wav2vec2', 'hubert', 'wavlm', 'zhao19', 'emo18'], required=True)
    parser_eval.add_argument('--target_model_name', type=str, choices=['wav2vec2', 'hubert', 'wavlm', 'zhao19', 'emo18'], required=True)
    parser_eval.add_argument('--rnn_name', type=str, default='lstm', choices=['gru', 'lstm'])
    parser_eval.add_argument('--target', type=int, default=-1, help='-1 if untargeted')
    parser_eval.add_argument('--instruments', type=str, nargs='+', default=["bass", "drums", "other", "vocals"])
    parser_eval.add_argument('--depth', type=int, default=1,help="Number of convs per block")
    parser_eval.add_argument('--features', type=int, default=32, help='Number of feature channels per layer')
    parser_eval.add_argument('--levels', type=int, default=6, help="Number of DS/US blocks")
    parser_eval.add_argument('--conv_type', type=str, default="gn",
                                help="Type of convolution (normal, BN-normalised, GN-normalised): normal/bn/gn")
    parser_eval.add_argument('--res', type=str, default="learned",
                                help="Resampling strategy: fixed sinc-based lowpass filtering or learned conv layer: fixed/learned")
    args = parser.parse_args()

    args.filename = get_filename(__file__)


    # Create log
    logs_dir = os.path.join(args.workspace, 'sparse_attack', args.filename, '0518', '{}_{}'.format(args.source_model_name, args.target_model_name))
    create_folder(logs_dir)
    custom = 'val={}'.format(args.validation)
    create_logging(logs_dir, custom, filemode='w')
    logging.info(args)


    if args.mode == 'eval':
        train(args)
    else:
        raise Exception('Error argument!')