voicebox/scripts/experiments/train.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchaudio
import psutil
import pickle
import librosa as li

from torch.utils.data import TensorDataset

import time
import random
import argparse
from datetime import datetime

import numpy as np
import pandas as pd
from typing import Dict
from pathlib import Path
from tqdm import tqdm
import builtins

from typing import Iterable
from copy import deepcopy

from distutils.util import strtobool

from src.data import *
from src.constants import *
from src.models import *
from src.simulation import *
from src.preprocess import *
from src.attacks.offline import *
from src.loss import *
from src.pipelines import *
from src.utils import *

################################################################################
# Train VoiceBox attack
################################################################################

BATCH_SIZE = 20     # training batch size
EPOCHS = 10         # training epochs
TARGET_PCTL = 25    # de-identification strength; in [1,5,10,15,20,25,50,90,100]
N_EMBEDDINGS_TRAIN = 15
TARGETED = False
TARGETS_TRAIN = 'centroid'  # 'random', 'same', 'single', 'median'
TARGETS_TEST = 'centroid'  # 'random', 'same', 'single', 'median'

# distributions of inter- ('targeted') and intra- ('untargeted') speaker
# distances in each pre-trained model's embedding spaces, as measured between
# individual utterances and their speaker centroid ('single-centroid') or
# between all pairs of individual utterances ('single-single') over the
# LibriSpeech test-clean dataset. This allows specification of attack strength
# during the training process
percentiles = {
    'resnet': {
        'targeted': {
            'single-centroid': {1:.495, 5:.572, 10:.617, 15:.648, 20:.673, 25:.695, 50:.773, 90:.892, 100:1.127},
            'single-single': {1:.560, 5:.630, 10:.672, 15:.700, 20:.722, 25:.742, 50:.813, 90:.924, 100:1.194}
        },
        'untargeted': {
            'single-centroid': {1:.099, 5:.117, 10:.126, 15:.133, 20:.139, 25:.145, 50:.170, 90:.253, 100:.587},
            'single-single': {1:.181, 5:.215, 10:.235, 15:.249, 20:.262, 25:.272, 50:.323, 90:.464, 100:.817}
        },
    },
    'yvector': {
        'targeted': {
            'single-centroid': {1:.665, 5:.757, 10:.801, 15:.830, 20:.851, 25:.868, 50:.936, 90:1.056, 100:1.312},
            'single-single': {1:.695, 5:.779, 10:.821, 15:.847, 20:.868, 25:.885, 50:.952, 90:1.072, 100:1.428}
        },
        'untargeted': {
            'single-single': {1:.218, 5:.268, 10:.301, 15:.325, 20:.345, 25:.365, 50:.455, 90:.684, 100:1.156},
            'single-centroid': {1:.114, 5:.143, 10:.159, 15:.170, 20:.180, 25:.190, 50:.242, 90:.413, 100:.874}
        }
    },
}


def set_random_seed(seed: int = 123):
    """Set random seed to allow for reproducibility"""
    random.seed(seed)
    torch.manual_seed(seed)

    if torch.backends.cudnn.is_available():
        # torch.backends.cudnn.benchmark = True
        torch.backends.cudnn.deterministic = True


def param_count(m: nn.Module, trainable: bool = False):
    """Count the number of trainable parameters (weights) in a model"""
    if trainable:
        return builtins.sum(
            [p.shape.numel() for p in m.parameters() if p.requires_grad])
    else:
        return builtins.sum([p.shape.numel() for p in m.parameters()])


def main():

    set_random_seed(0)

    model = SpeakerVerificationModel(
        model=ResNetSE34V2(nOut=512, encoder_type='ASP'),
        n_segments=1,
        segment_select='lin',
        distance_fn='cosine',
        threshold=percentiles['resnet']['targeted']['single-centroid' if
        TARGETS_TRAIN == 'centroid' else 'single-single'][TARGET_PCTL]
    )
    model.load_weights(MODELS_DIR / 'speaker' / 'resnetse34v2' / 'resnetse34v2.pt')

    # instantiate training pipeline
    pipeline = Pipeline(
        simulation=None,
        preprocessor=Preprocessor(Normalize(method='peak')),
        model=model,
        device='cuda' if torch.cuda.is_available() else 'cpu'
    )

    attacks = {}

    # log training progress
    writer = Writer(
        root_dir=RUNS_DIR,
        name='train-attacks',
        use_timestamp=True,
        log_iter=300,
        use_tb=True
    )

    # adversarial training loss
    adv_loss = SpeakerEmbeddingLoss(
        targeted=TARGETED,
        confidence=0.1,
        threshold=pipeline.model.threshold
    )

    # auxiliary loss
    aux_loss = SumLoss().add_loss_function(
        DemucsMRSTFTLoss(), 1.0
    ).add_loss_function(L1Loss(), 1.0).to('cuda')

    # speech features loss actually seems to do better...
    # aux_loss = SumLoss().add_loss_function(SpeechFeatureLoss(), 1e-6).to('cuda')

    attacks['voicebox'] = VoiceBoxAttack(
        pipeline=pipeline,
        adv_loss=adv_loss,
        aux_loss=aux_loss,
        lr=1e-4,
        epochs=EPOCHS,
        batch_size=BATCH_SIZE,
        voicebox_kwargs={
            'win_length': 256,
            'ppg_encoder_hidden_size': 256,
            'use_phoneme_encoder': True,
            'use_pitch_encoder': True,
            'use_loudness_encoder': True,
            'spec_encoder_lookahead_frames': 0,
            'spec_encoder_type': 'mel',
            'spec_encoder_mlp_depth': 2,
            'bottleneck_lookahead_frames': 5,
            'ppg_encoder_path': PPG_PRETRAINED_PATH,
            'n_bands': 128,
            'spec_encoder_hidden_size': 512,
            'bottleneck_skip': True,
            'bottleneck_hidden_size': 512,
            'bottleneck_feedforward_size': 512,
            'bottleneck_type': 'lstm',
            'bottleneck_depth': 2,
            'control_eps': 0.5,
            'projection_norm': float('inf'),
            'conditioning_dim': 512
        },
        writer=writer,
        checkpoint_name='voicebox-attack'
    )

    attacks['universal'] = AdvPulseAttack(
        pipeline=pipeline,
        adv_loss=adv_loss,
        pgd_norm=float('inf'),
        pgd_variant=None,
        scale_grad=None,
        eps=0.08,
        length=2.0,
        align='random',  # 'start',
        lr=1e-4,
        normalize=True,
        loop=True,
        aux_loss=aux_loss,
        epochs=EPOCHS,
        batch_size=BATCH_SIZE,
        writer=writer,
        checkpoint_name='universal-attack'
    )

    if torch.cuda.is_available():

        # prepare for multi-GPU training
        device_ids = get_cuda_device_ids()

        # wrap pipeline for multi-GPU training
        pipeline = wrap_pipeline_multi_gpu(pipeline, device_ids)

    # load training and validation datasets. Features will be computed and
    # cached to disk, which may take some time
    data_train = LibriSpeechDataset(
        split='train-clean-100', features=['pitch', 'periodicity', 'loudness'])
    data_test = LibriSpeechDataset(
        split='test-clean', features=['pitch', 'periodicity', 'loudness'])

    # reassign targets if necessary
    compiled_train, compiled_test = create_embedding_dataset(
        pipeline=pipeline,
        select_train=TARGETS_TRAIN,
        select_test=TARGETS_TEST,
        data_train=data_train,
        data_test=data_test,
        targeted=TARGETED,
        target_class=None,
        num_embeddings_train=N_EMBEDDINGS_TRAIN,
        batch_size=20
    )

    # extract embedding datasets
    data_train = compiled_train['dataset']
    data_test = compiled_test['dataset']

    # log memory use prior to training
    writer.log_info(f'Training data ready; memory use: '
                    f'{psutil.virtual_memory().percent :0.3f}%')
    writer.log_cuda_memory()

    for attack_name, attack in attacks.items():

        writer.log_info(f'Preparing {attack_name}...')

        if torch.cuda.is_available():

            attack.perturbation.to('cuda')
            attack.pipeline.to('cuda')

            # wrap attack for multi-GPU training
            attack = wrap_attack_multi_gpu(attack, device_ids)

        # evaluate performance
        with torch.no_grad():
            x_example = next(iter(data_train))['x'].to(pipeline.device)
            st = time.time()
            outs = attack.perturbation(x_example)
            dur = time.time() - st

        writer.log_info(
            f'Processing time per input (device: '
            f'{pipeline.device}): {dur/x_example.shape[0] :0.4f} (s)'
        )
        writer.log_info(f'Trainable parameters: '
                        f'{param_count(attack.perturbation, trainable=True)}')
        writer.log_info(f'Total parameters: {param_count(attack.perturbation, trainable=False)}')

        # train
        writer.log_info('Training attack...')
        attack.train(data_train=data_train, data_val=data_test)

        # evaluate
        writer.log_info(f'Evaluating attack...')
        x_adv, success, detection = attack.evaluate(
            dataset=data_test
        )

        # log results summary: success rate in achieving target threshold
        writer.log_info(
            f'Success rate in meeting embedding distance threshold {pipeline.model.threshold}'
            f' ({TARGET_PCTL}%): '
            f'{success.flatten().mean().item()}'
        )


if __name__ == "__main__":
    main()