adversarial.py

# Adversarial learning for event-based music generation with SeqGAN
# Reference:
# "SeqGAN: Sequence Generative Adversarial Nets with Policy Gradient."
# (Yu, Lantao, et al.).
# ... Honestly, it's too hard to train ;(

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.distributions import Categorical

import numpy as np
import os, sys, time, argparse
from progress.bar import Bar

import config, utils
from config import device
from data import Dataset
from model import PerformanceRNN
from sequence import EventSeq, ControlSeq

# pylint: disable=E1101


#========================================================================
# Discriminator
#========================================================================

discriminator_config = {
    'event_dim': EventSeq.dim(),
    'hidden_dim': 512,
    'gru_layers': 3,
    'gru_dropout': 0.3
}

class EventSequenceEncoder(nn.Module):
    def __init__(self, event_dim=EventSeq.dim(), hidden_dim=512,
                 gru_layers=3, gru_dropout=0.3):
        super().__init__()
        self.event_embedding = nn.Embedding(event_dim, hidden_dim)
        self.gru = nn.GRU(hidden_dim, hidden_dim,
                          num_layers=gru_layers, dropout=gru_dropout)
        self.attn = nn.Parameter(torch.randn(hidden_dim), requires_grad=True)
        self.output_fc = nn.Linear(hidden_dim, 1)
        self.output_fc_activation = nn.Sigmoid()

    def forward(self, events, hidden=None, output_logits=False):
        # events: [steps, batch_size]
        events = self.event_embedding(events)
        outputs, _ = self.gru(events, hidden) # [t, b, h]
        weights = (outputs * self.attn).sum(-1, keepdim=True)
        output = (outputs * weights).mean(0) # [b, h]
        output = self.output_fc(output).squeeze(-1) # [b]
        if output_logits:
            return output
        output = self.output_fc_activation(output)
        return output


#========================================================================
# Pretrain Discriminator
#========================================================================

def pretrain_discriminator(model_sess_path,         # load
                           discriminator_sess_path, # load + save
                           batch_data_generator,    # Dataset(...).batches(...)
                           discriminator_config_overwrite={},
                           gradient_clipping=False,
                           control_ratio=1.0,
                           num_iter=-1,
                           save_interval=60.0,
                           discriminator_lr=0.001,
                           enable_logging=False,
                           auto_sample_factor=False,
                           sample_factor=1.0):

    print('-' * 70)
    print('model_sess_path:', model_sess_path)
    print('discriminator_sess_path:', discriminator_sess_path)
    print('discriminator_config_overwrite:', discriminator_config_overwrite)
    print('sample_factor:', sample_factor)
    print('auto_sample_factor:', auto_sample_factor)
    print('discriminator_lr:', discriminator_lr)
    print('gradient_clipping:', gradient_clipping)
    print('control_ratio:', control_ratio)
    print('num_iter:', num_iter)
    print('save_interval:', save_interval)
    print('enable_logging:', enable_logging)
    print('-' * 70)
    
    # Load generator
    model_sess = torch.load(model_sess_path)
    model_config = model_sess['model_config']
    model = PerformanceRNN(**model_config).to(device)
    model.load_state_dict(model_sess['model_state'])

    print(f'Generator from "{model_sess_path}"')
    print(model)
    print('-' * 70)

    # Load discriminator and optimizer
    global discriminator_config
    try:
        discriminator_sess = torch.load(discriminator_sess_path)
        discriminator_config = discriminator_sess['discriminator_config']
        discriminator_state = discriminator_sess['discriminator_state']
        discriminator_optimizer_state = discriminator_sess['discriminator_optimizer_state']
        print(f'Discriminator from "{discriminator_sess_path}"')
        discriminator_loaded = True
    except:
        print(f'New discriminator session at "{discriminator_sess_path}"')
        discriminator_config.update(discriminator_config_overwrite)
        discriminator_loaded = False

    discriminator = EventSequenceEncoder(**discriminator_config).to(device)
    optimizer = optim.Adam(discriminator.parameters(), lr=discriminator_lr)
    if discriminator_loaded:
        discriminator.load_state_dict(discriminator_state)
        optimizer.load_state_dict(discriminator_optimizer_state)

    print(discriminator)
    print(optimizer)
    print('-' * 70)

    def save_discriminator():
        print(f'Saving to "{discriminator_sess_path}"')
        torch.save({
            'discriminator_config': discriminator_config,
            'discriminator_state': discriminator.state_dict(),
            'discriminator_optimizer_state': optimizer.state_dict()
        }, discriminator_sess_path)
        print('Done saving')

    # Disable gradient for generator
    for parameter in model.parameters():
        parameter.requires_grad_(False)

    model.eval()
    discriminator.train()

    loss_func = nn.BCEWithLogitsLoss()
    last_save_time = time.time()

    if enable_logging:
        from tensorboardX import SummaryWriter
        writer = SummaryWriter()

    try:
        for i, (events, controls) in enumerate(batch_data_generator):
            if i == num_iter:
                break
            
            steps, batch_size = events.shape

            # Prepare inputs
            events = torch.LongTensor(events).to(device)
            if np.random.random() <= control_ratio:
                controls = torch.FloatTensor(controls).to(device)
            else:
                controls = None

            init = torch.randn(batch_size, model.init_dim).to(device)

            # Predict for real event sequence
            real_events = events
            real_logit = discriminator(real_events, output_logits=True)
            real_target = torch.ones_like(real_logit).to(device)

            if auto_sample_factor:
                sample_factor = np.random.choice([
                    0.1, 0.4, 0.6, 0.7, 0.8, 0.9, 1.0,
                    1.1, 1.2, 1.4, 1.6, 2.0, 4.0, 10.0])

            # Predict for fake event sequence from the generator
            fake_events = model.generate(init, steps, None, controls,
                                         greedy=0, output_type='index',
                                         temperature=sample_factor)
            fake_logit = discriminator(fake_events, output_logits=True)
            fake_target = torch.zeros_like(fake_logit).to(device)

            # Compute loss
            loss = (loss_func(real_logit, real_target) +
                    loss_func(fake_logit, fake_target)) / 2
            
            # Backprop
            discriminator.zero_grad()
            loss.backward()

            # Gradient clipping
            norm = utils.compute_gradient_norm(discriminator.parameters())
            if gradient_clipping:
                nn.utils.clip_grad_norm_(discriminator.parameters(), gradient_clipping)

            optimizer.step()

            # Logging
            loss = loss.item()
            norm = norm.item()
            print(f'{i} loss: {loss}, norm: {norm}, sf: {sample_factor}')
            if enable_logging:
                writer.add_scalar(f'pretrain/D/loss/all', loss, i)
                writer.add_scalar(f'pretrain/D/loss/{sample_factor}', loss, i)
                writer.add_scalar(f'pretrain/D/norm/{sample_factor}', norm, i)

            if last_save_time + save_interval < time.time():
                last_save_time = time.time()
                save_discriminator()

    except KeyboardInterrupt:
        save_discriminator()


#========================================================================
# Adversarial Learning
#========================================================================


def train_adversarial(sess_path, batch_data_generator,
                      model_load_path, model_optimizer_class,
                      model_gradient_clipping, discriminator_gradient_clipping,
                      model_learning_rate, reset_model_optimizer,
                      discriminator_load_path, discriminator_optimizer_class,
                      discriminator_learning_rate, reset_discriminator_optimizer,
                      g_max_q_mean, g_min_q_mean, d_min_loss, g_max_steps, d_max_steps,
                      mc_sample_size, mc_sample_factor, first_to_train,
                      save_interval, control_ratio, enable_logging):
    
    if enable_logging:
        from tensorboardX import SummaryWriter
        writer = SummaryWriter()

    if os.path.isfile(sess_path):
        adv_state = torch.load(sess_path)
        model_config = adv_state['model_config']
        model_state = adv_state['model_state']
        model_optimizer_state = adv_state['model_optimizer_state']
        discriminator_config = adv_state['discriminator_config']
        discriminator_state = adv_state['discriminator_state']
        discriminator_optimizer_state = adv_state['discriminator_optimizer_state']
        print('-' * 70)
        print('Session is loaded from', sess_path)
        loaded_from_session = True

    else:
        model_sess = torch.load(model_load_path)
        model_config = model_sess['model_config']
        model_state = model_sess['model_state']
        discriminator_sess = torch.load(discriminator_load_path)
        discriminator_config = discriminator_sess['discriminator_config']
        discriminator_state = discriminator_sess['discriminator_state']
        loaded_from_session = False

    model = PerformanceRNN(**model_config)
    model.load_state_dict(model_state)
    model.to(device).train()
    model_optimizer = model_optimizer_class(model.parameters(), lr=model_learning_rate)

    discriminator = EventSequenceEncoder(**discriminator_config)
    discriminator.load_state_dict(discriminator_state)
    discriminator.to(device).train()
    discriminator_optimizer = discriminator_optimizer_class(discriminator.parameters(),
                                                            lr=discriminator_learning_rate)

    if loaded_from_session:
        if not reset_model_optimizer:
            model_optimizer.load_state_dict(model_optimizer_state)
        if not reset_discriminator_optimizer:
            discriminator_optimizer.load_state_dict(discriminator_optimizer_state)

    g_loss_func = nn.CrossEntropyLoss()
    d_loss_func = nn.BCEWithLogitsLoss(reduce=False)
    

    print('-' * 70)
    print('Options')
    print('sess_path:', sess_path)
    print('save_interval:', save_interval)
    print('batch_data_generator:', batch_data_generator)
    print('control_ratio:', control_ratio)
    print('g_max_q_mean:', g_max_q_mean)
    print('g_min_q_mean:', g_min_q_mean)
    print('d_min_loss:', d_min_loss)
    print('mc_sample_size:', mc_sample_size)
    print('mc_sample_factor:', mc_sample_factor)
    print('enable_logging:', enable_logging)
    print('model_load_path:', model_load_path)
    print('model_loss:', g_loss_func)
    print('model_optimizer_class:', model_optimizer_class)
    print('model_gradient_clipping:', model_gradient_clipping)
    print('model_learning_rate:', model_learning_rate)
    print('reset_model_optimizer:', reset_model_optimizer)
    print('discriminator_load_path:', discriminator_load_path)
    print('discriminator_loss:', d_loss_func)
    print('discriminator_optimizer_class:', discriminator_optimizer_class)
    print('discriminator_gradient_clipping:', discriminator_gradient_clipping)
    print('discriminator_learning_rate:', discriminator_learning_rate)
    print('reset_discriminator_optimizer:', reset_discriminator_optimizer)
    print('first_to_train:', first_to_train)
    print('-' * 70)
    print(f'Generator from "{sess_path if loaded_from_session else model_load_path}"')
    print(model)
    print(model_optimizer)
    print('-' * 70)
    print(f'Discriminator from "{sess_path if loaded_from_session else discriminator_load_path}"')
    print(discriminator)
    print(discriminator_optimizer)
    print('-' * 70)
    
    
    def save():
        print(f'Saving to "{sess_path}"')
        torch.save({
            'model_config': model_config,
            'model_state': model.state_dict(),
            'model_optimizer_state': model_optimizer.state_dict(),
            'discriminator_config': discriminator_config,
            'discriminator_state': discriminator.state_dict(),
            'discriminator_optimizer_state': discriminator_optimizer.state_dict()
        }, sess_path)
        print('Done saving')
    
    def mc_rollout(generated, hidden, total_steps, controls=None):
        # generated: [t, batch_size]
        # hidden: [n_layers, batch_size, hidden_dim]
        # controls: [total_steps - t, batch_size, control_dim]
        generated = torch.cat(generated, 0)
        generated_steps, batch_size = generated.shape # t, b
        steps = total_steps - generated_steps # s

        generated = generated.unsqueeze(1) # [t, 1, b]
        generated = generated.repeat(1, mc_sample_size, 1) # [t, mcs, b]
        generated = generated.view(generated_steps, -1) # [t, mcs * b]

        hidden = hidden.unsqueeze(1).repeat(1, mc_sample_size, 1, 1)
        hidden = hidden.view(model.gru_layers, -1, model.hidden_dim)

        if controls is not None:
            assert controls.shape == (steps, batch_size, model.control_dim)
            controls = controls.unsqueeze(1) # [s, 1, b, c]
            controls = controls.repeat(1, mc_sample_size, 1, 1) # [s, mcs, b, c]
            controls = controls.view(steps, -1, model.control_dim) # [s, mcs * b, c]

        event = generated[-1].unsqueeze(0) # [1, mcs * b]
        control = None # default when controls is None
        outputs = []

        for i in range(steps):
            if controls is not None:
                control = controls[i].unsqueeze(0) # [1, mcs * b, c]

            output, hidden = model.forward(event, control=control, hidden=hidden)
            probs = model.output_fc_activation(output / mc_sample_factor)
            event = Categorical(probs).sample() # [1, mcs * b]
            outputs.append(event)

        sequences = torch.cat([generated, *outputs], 0)
        assert sequences.shape == (total_steps, mc_sample_size * batch_size)
        return sequences


    def train_generator(batch_size, init, events, controls):
        # Generator step
        hidden = model.init_to_hidden(init)
        event = model.get_primary_event(batch_size)
        outputs = []
        generated = []
        q_values = []

        for step in Bar('MC Rollout').iter(range(steps)):
            control = controls[step].unsqueeze(0) if use_control else None
            output, hidden = model.forward(event, control=control, hidden=hidden)
            outputs.append(output)
            probs = model.output_fc_activation(output / mc_sample_factor)
            generated.append(Categorical(probs).sample())

            with torch.no_grad():
                if step < steps - 1:
                    sequences = mc_rollout(generated, hidden, steps, controls[step+1:])
                    mc_score = discriminator(sequences) # [mcs * b]
                    mc_score = mc_score.view(mc_sample_size, batch_size) # [mcs, b]
                    q_value = mc_score.mean(0, keepdim=True) # [1, batch_size]
                
                else:
                    q_value = discriminator(torch.cat(generated, 0))
                    q_value = q_value.unsqueeze(0) # [1, batch_size]
            
                q_values.append(q_value)
        
        # Compute loss
        q_values = torch.cat(q_values, 0) # [steps, batch_size]
        q_mean = q_values.mean().detach()
        q_values = q_values - q_mean
        generated = torch.cat(generated, 0) # [steps, batch_size]
        outputs = torch.cat(outputs, 0) # [steps, batch_size, event_dim]
        loss = F.cross_entropy(outputs.view(-1, model.event_dim),
                               generated.view(-1),
                               reduce=False)
        loss = (loss * q_values.view(-1)).mean()

        # Backprop
        model.zero_grad()
        loss.backward()

        # Gradient clipping
        norm = utils.compute_gradient_norm(model.parameters())
        if model_gradient_clipping:
            nn.utils.clip_grad_norm_(model.parameters(), model_gradient_clipping)

        model_optimizer.step()

        q_mean = q_mean.item()
        norm = norm.item()
        return q_mean, norm

    def train_discriminator(batch_size, init, events, controls):
        # Discriminator step
        with torch.no_grad():
            generated = model.generate(init, steps, None, controls,
                                        greedy=0, temperature=mc_sample_factor)
            
        fake_logit = discriminator(generated, output_logits=True)
        real_logit = discriminator(events, output_logits=True)
        fake_target = torch.zeros_like(fake_logit)
        real_target = torch.ones_like(real_logit)

        # Compute loss
        fake_loss = F.binary_cross_entropy_with_logits(fake_logit, fake_target)
        real_loss = F.binary_cross_entropy_with_logits(real_logit, real_target)
        loss = (real_loss + fake_loss) / 2

        # Backprop
        discriminator.zero_grad()
        loss.backward()

        # Gradient clipping
        norm = utils.compute_gradient_norm(discriminator.parameters())
        if discriminator_gradient_clipping:
            nn.utils.clip_grad_norm_(discriminator.parameters(), discriminator_gradient_clipping)

        discriminator_optimizer.step()

        real_loss = real_loss.item()
        fake_loss = fake_loss.item()
        loss = loss.item()
        norm = norm.item()
        return loss, real_loss, fake_loss, norm

    try:
        last_save_time = time.time()
        step_for = first_to_train
        g_steps = 0
        d_steps = 0
        
        for i, (events, controls) in enumerate(batch_data_generator):
            steps, batch_size = events.shape
            init = torch.randn(batch_size, model.init_dim).to(device)
            events = torch.LongTensor(events).to(device)

            use_control = np.random.random() <= control_ratio
            controls = torch.FloatTensor(controls).to(device) if use_control else None

            if step_for == 'G':
                q_mean, norm = train_generator(batch_size, init, events, controls)
                g_steps += 1

                print(f'{i} (G-step) Q_mean: {q_mean}, norm: {norm}')
                if enable_logging:
                    writer.add_scalar('adversarial/G/Q_mean', q_mean, i)
                    writer.add_scalar('adversarial/G/norm', norm, i)

                if q_mean < g_min_q_mean:
                    print(f'Q is too small: {q_mean}, exiting')
                    raise KeyboardInterrupt

                if q_mean > g_max_q_mean or (g_max_steps and g_steps >= g_max_steps):
                    step_for = 'D'
                    d_steps = 0

            if step_for == 'D':
                loss, real_loss, fake_loss, norm = train_discriminator(batch_size, init, events, controls)
                d_steps += 1

                print(f'{i} (D-step) loss: {loss} (real: {real_loss}, fake: {fake_loss}), norm: {norm}')
                if enable_logging:
                    writer.add_scalar('adversarial/D/loss', loss, i)
                    writer.add_scalar('adversarial/D/norm', norm, i)

                if fake_loss <= real_loss < d_min_loss or (d_max_steps and d_steps >= d_max_steps):
                    step_for = 'G'
                    g_steps = 0

            if last_save_time + save_interval < time.time():
                last_save_time = time.time()
                save()

    except KeyboardInterrupt:
        save()



#========================================================================
# Script Arguments
#========================================================================

def batch_generator(args):
    print('-' * 70)
    dataset = Dataset(args.dataset_path, verbose=True)
    print(dataset)
    return dataset.batches(args.batch_size, args.window_size, args.stride_size)

def pretrain(args):
    pretrain_discriminator(model_sess_path=args.generator_session_path,
                           discriminator_sess_path=args.discriminator_session_path,
                           discriminator_config_overwrite=utils.params2dict(args.discriminator_parameters),
                           batch_data_generator=args.batch_generator(args),
                           gradient_clipping=args.gradient_clipping,
                           sample_factor=args.sample_factor,
                           auto_sample_factor=args.auto_sample_factor,
                           control_ratio=args.control_ratio,
                           num_iter=args.stop_iteration,
                           save_interval=args.save_interval,
                           discriminator_lr=args.discriminator_learning_rate,
                           enable_logging=args.enable_logging)

def adversarial(args):
    train_adversarial(sess_path=args.session_path,
                      batch_data_generator=args.batch_generator(args),
                      model_load_path=args.generator_load_path,
                      discriminator_load_path=args.discriminator_load_path,
                      model_optimizer_class=getattr(optim, args.generator_optimizer),
                      discriminator_optimizer_class=getattr(optim, args.discriminator_optimizer),
                      model_gradient_clipping=args.generator_gradient_clipping,
                      discriminator_gradient_clipping=args.discriminator_gradient_clipping,
                      model_learning_rate=args.generator_learning_rate,
                      discriminator_learning_rate=args.discriminator_learning_rate,
                      reset_model_optimizer=args.reset_generator_optimizer,
                      reset_discriminator_optimizer=args.reset_discriminator_optimizer,
                      g_max_q_mean=args.g_max_q_mean,
                      g_min_q_mean=args.g_min_q_mean,
                      d_min_loss=args.d_min_loss,
                      g_max_steps=args.g_max_steps,
                      d_max_steps=args.d_max_steps,
                      mc_sample_size=args.monte_carlo_sample_size,
                      mc_sample_factor=args.monte_carlo_sample_factor,
                      first_to_train=args.first_to_train,
                      control_ratio=args.control_ratio,
                      save_interval=args.save_interval,
                      enable_logging=args.enable_logging)

def get_args():
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers()
    parser.add_argument('-d', '--dataset-path', type=str, required=True)
    parser.add_argument('-b', '--batch-size', type=int, default=64)
    parser.add_argument('-w', '--window-size', type=int, default=200)
    parser.add_argument('-s', '--stride-size', type=int, default=10)
    parser.set_defaults(batch_generator=batch_generator)
    pre_parser = subparsers.add_parser('pretrain', aliases=['p', 'pre'])
    pre_parser.add_argument('-G', '--generator-session-path', type=str, default=True)
    pre_parser.add_argument('-D', '--discriminator-session-path', type=str, required=True)
    pre_parser.add_argument('-p', '--discriminator-parameters', type=str, default='')
    pre_parser.add_argument('-l', '--discriminator-learning-rate', type=float, default=0.001)
    pre_parser.add_argument('-g', '--gradient-clipping', type=float, default=False)
    pre_parser.add_argument('-f', '--sample-factor', type=float, default=1.0)
    pre_parser.add_argument('-af', '--auto-sample-factor', action='store_true', default=False)
    pre_parser.add_argument('-c', '--control-ratio', type=float, default=1.0)
    pre_parser.add_argument('-n', '--stop-iteration', type=int, default=-1)
    pre_parser.add_argument('-i', '--save-interval', type=float, default=60.0)
    pre_parser.add_argument('-L', '--enable-logging', action='store_true', default=False)
    pre_parser.set_defaults(main=pretrain)
    adv_parser = subparsers.add_parser('adversarial', aliases=['a', 'adv'])
    adv_parser.add_argument('-S', '--session-path', type=str, required=True)
    adv_parser.add_argument('-Gp', '--generator-load-path', type=str)
    adv_parser.add_argument('-Dp', '--discriminator-load-path', type=str)
    adv_parser.add_argument('-Go', '--generator-optimizer', type=str, default='Adam')
    adv_parser.add_argument('-Do', '--discriminator-optimizer', type=str, default='RMSprop')
    adv_parser.add_argument('-Gg', '--generator-gradient-clipping', type=float, default=False)
    adv_parser.add_argument('-Dg', '--discriminator-gradient-clipping', type=float, default=False)
    adv_parser.add_argument('-Gl', '--generator-learning-rate', type=float, default=0.001)
    adv_parser.add_argument('-Dl', '--discriminator-learning-rate', type=float, default=0.001)
    adv_parser.add_argument('-Gr', '--reset-generator-optimizer', action='store_true', default=False)
    adv_parser.add_argument('-Dr', '--reset-discriminator-optimizer', action='store_true', default=False)
    adv_parser.add_argument('-Gq', '--g-max-q-mean', type=float, default=0.5)
    adv_parser.add_argument('-Gm', '--g-min-q-mean', type=float, default=0.0)
    adv_parser.add_argument('-Dm', '--d-min-loss', type=float, default=0.5)
    adv_parser.add_argument('-Gs', '--g-max-steps', type=int, default=0)
    adv_parser.add_argument('-Ds', '--d-max-steps', type=int, default=0)
    adv_parser.add_argument('-f', '--first-to-train', type=str, default='G', choices=['G', 'D'])
    adv_parser.add_argument('-ms', '--monte-carlo-sample-size', type=int, default=8)
    adv_parser.add_argument('-mf', '--monte-carlo-sample-factor', type=float, default=1.0)
    adv_parser.add_argument('-c', '--control-ratio', type=float, default=1.0)
    adv_parser.add_argument('-i', '--save-interval', type=float, default=60.0)
    adv_parser.add_argument('-L', '--enable-logging', action='store_true', default=False)
    adv_parser.set_defaults(main=adversarial)
    return parser.parse_args()


if __name__ == '__main__':
    args = get_args()
    args.main(args)