trainer.py

from __future__ import print_function

import argparse
import torch

import numpy as np

from torch.autograd.variable import Variable
from torch.optim import Adam

from multi_categorical_gans.datasets.dataset import Dataset
from multi_categorical_gans.datasets.formats import data_formats, loaders

from multi_categorical_gans.methods.general.autoencoder import AutoEncoder
from multi_categorical_gans.methods.general.generator import Generator
from multi_categorical_gans.methods.general.discriminator import Discriminator
from multi_categorical_gans.methods.general.wgan_gp import calculate_gradient_penalty

from multi_categorical_gans.utils.categorical import load_variable_sizes_from_metadata, categorical_variable_loss
from multi_categorical_gans.utils.commandline import DelayedKeyboardInterrupt, parse_int_list
from multi_categorical_gans.utils.cuda import to_cuda_if_available, to_cpu_if_available
from multi_categorical_gans.utils.initialization import load_or_initialize
from multi_categorical_gans.utils.logger import Logger


def add_noise_to_code(code, noise_radius):
    if noise_radius > 0:
        means = torch.zeros_like(code)
        gauss_noise = torch.normal(means, noise_radius)
        return code + to_cuda_if_available(Variable(gauss_noise))
    else:
        return code


def train(autoencoder,
          generator,
          discriminator,
          train_data,
          val_data,
          output_ae_path,
          output_gen_path,
          output_disc_path,
          output_loss_path,
          batch_size=1000,
          start_epoch=0,
          num_epochs=1000,
          num_ae_steps=1,
          num_disc_steps=2,
          num_gen_steps=1,
          noise_size=128,
          l2_regularization=0.001,
          learning_rate=0.001,
          ae_noise_radius=0.2,
          ae_noise_anneal=0.995,
          normalize_code=True,
          variable_sizes=None,
          temperature=None,
          penalty=0.1
          ):
    autoencoder, generator, discriminator = to_cuda_if_available(autoencoder, generator, discriminator)

    optim_ae = Adam(autoencoder.parameters(), weight_decay=l2_regularization, lr=learning_rate)
    optim_gen = Adam(generator.parameters(), weight_decay=l2_regularization, lr=learning_rate)
    optim_disc = Adam(discriminator.parameters(), weight_decay=l2_regularization, lr=learning_rate)

    logger = Logger(output_loss_path, append=start_epoch > 0)

    for epoch_index in range(start_epoch, num_epochs):
        logger.start_timer()

        # train
        autoencoder.train(mode=True)
        generator.train(mode=True)
        discriminator.train(mode=True)

        ae_losses = []
        disc_losses = []
        gen_losses = []

        more_batches = True
        train_data_iterator = train_data.batch_iterator(batch_size)

        while more_batches:
            # train autoencoder
            for _ in range(num_ae_steps):
                try:
                    batch = next(train_data_iterator)
                except StopIteration:
                    more_batches = False
                    break

                autoencoder.zero_grad()

                batch_original = Variable(torch.from_numpy(batch))
                batch_original = to_cuda_if_available(batch_original)
                batch_code = autoencoder.encode(batch_original, normalize_code=normalize_code)
                batch_code = add_noise_to_code(batch_code, ae_noise_radius)

                batch_reconstructed = autoencoder.decode(batch_code, training=True, temperature=temperature)

                ae_loss = categorical_variable_loss(batch_reconstructed, batch_original, variable_sizes)
                ae_loss.backward()

                optim_ae.step()

                ae_loss = to_cpu_if_available(ae_loss)
                ae_losses.append(ae_loss.data.numpy())

            # train discriminator
            for _ in range(num_disc_steps):
                try:
                    batch = next(train_data_iterator)
                except StopIteration:
                    more_batches = False
                    break

                discriminator.zero_grad()
                autoencoder.zero_grad()

                # first train the discriminator only with real data
                real_features = Variable(torch.from_numpy(batch))
                real_features = to_cuda_if_available(real_features)
                real_code = autoencoder.encode(real_features, normalize_code=normalize_code)
                real_code = add_noise_to_code(real_code, ae_noise_radius)
                real_pred = discriminator(real_code)
                real_loss = - real_pred.mean(0).view(1)
                real_loss.backward()

                # then train the discriminator only with fake data
                noise = Variable(torch.FloatTensor(len(batch), noise_size).normal_())
                noise = to_cuda_if_available(noise)
                fake_code = generator(noise)
                fake_code = fake_code.detach()  # do not propagate to the generator
                fake_pred = discriminator(fake_code)
                fake_loss = fake_pred.mean(0).view(1)
                fake_loss.backward()

                # this is the magic from WGAN-GP
                gradient_penalty = calculate_gradient_penalty(discriminator, penalty, real_code, fake_code)
                gradient_penalty.backward()

                optim_ae.step()
                optim_disc.step()

                disc_loss = real_loss + fake_loss + gradient_penalty
                disc_loss = to_cpu_if_available(disc_loss)
                disc_losses.append(disc_loss.data.numpy())

                del disc_loss
                del gradient_penalty
                del fake_loss
                del real_loss

            # train generator
            for _ in range(num_gen_steps):
                generator.zero_grad()

                noise = Variable(torch.FloatTensor(len(batch), noise_size).normal_())
                noise = to_cuda_if_available(noise)
                gen_code = generator(noise)
                fake_pred = discriminator(gen_code)
                fake_loss = - fake_pred.mean(0).view(1)
                fake_loss.backward()

                optim_gen.step()

                fake_loss = to_cpu_if_available(fake_loss)
                gen_losses.append(fake_loss.data.numpy()[0])

                del fake_loss

        # log epoch metrics for current class
        logger.log(epoch_index, num_epochs, "autoencoder", "train_mean_loss", np.mean(ae_losses))
        logger.log(epoch_index, num_epochs, "discriminator", "train_mean_loss", np.mean(disc_losses))
        logger.log(epoch_index, num_epochs, "generator", "train_mean_loss", np.mean(gen_losses))

        # save models for the epoch
        with DelayedKeyboardInterrupt():
            torch.save(autoencoder.state_dict(), output_ae_path)
            torch.save(generator.state_dict(), output_gen_path)
            torch.save(discriminator.state_dict(), output_disc_path)
            logger.flush()

        ae_noise_radius *= ae_noise_anneal

    logger.close()


def main():
    options_parser = argparse.ArgumentParser(description="Train ARAE or MC-ARAE. "
                                                         + "Define 'metadata' and 'temperature' to use MC-ARAE.")

    options_parser.add_argument("data", type=str, help="Training data. See 'data_format' parameter.")

    options_parser.add_argument("output_autoencoder", type=str, help="Autoencoder output file.")
    options_parser.add_argument("output_generator", type=str, help="Generator output file.")
    options_parser.add_argument("output_discriminator", type=str, help="Discriminator output file.")
    options_parser.add_argument("output_loss", type=str, help="Loss output file.")

    options_parser.add_argument("--input_autoencoder", type=str, help="Autoencoder input file.", default=None)
    options_parser.add_argument("--input_generator", type=str, help="Generator input file.", default=None)
    options_parser.add_argument("--input_discriminator", type=str, help="Discriminator input file.", default=None)

    options_parser.add_argument("--metadata", type=str,
                                help="Information about the categorical variables in json format.")

    options_parser.add_argument(
        "--validation_proportion", type=float,
        default=.1,
        help="Ratio of data for validation."
    )

    options_parser.add_argument(
        "--data_format",
        type=str,
        default="sparse",
        choices=data_formats,
        help="Either a dense numpy array or a sparse csr matrix."
    )

    options_parser.add_argument(
        "--code_size",
        type=int,
        default=128,
        help="Dimension of the autoencoder latent space."
    )

    options_parser.add_argument(
        "--noise_size",
        type=int,
        default=128,
        help="Dimension of the generator input noise."
    )

    options_parser.add_argument(
        "--encoder_hidden_sizes",
        type=str,
        default="",
        help="Size of each hidden layer in the encoder separated by commas (no spaces)."
    )

    options_parser.add_argument(
        "--decoder_hidden_sizes",
        type=str,
        default="",
        help="Size of each hidden layer in the decoder separated by commas (no spaces)."
    )

    options_parser.add_argument(
        "--batch_size",
        type=int,
        default=100,
        help="Amount of samples per batch."
    )

    options_parser.add_argument(
        "--start_epoch",
        type=int,
        default=0,
        help="Starting epoch."
    )

    options_parser.add_argument(
        "--num_epochs",
        type=int,
        default=5000,
        help="Number of epochs."
    )

    options_parser.add_argument(
        "--l2_regularization",
        type=float,
        default=0,
        help="L2 regularization weight for every parameter."
    )

    options_parser.add_argument(
        "--learning_rate",
        type=float,
        default=1e-5,
        help="Adam learning rate."
    )

    options_parser.add_argument(
        "--generator_hidden_sizes",
        type=str,
        default="100,100,100",
        help="Size of each hidden layer in the generator separated by commas (no spaces)."
    )

    options_parser.add_argument(
        "--bn_decay",
        type=float,
        default=0.9,
        help="Batch normalization decay for the generator and discriminator."
    )

    options_parser.add_argument(
        "--discriminator_hidden_sizes",
        type=str,
        default="100",
        help="Size of each hidden layer in the discriminator separated by commas (no spaces)."
    )

    options_parser.add_argument(
        "--num_autoencoder_steps",
        type=int,
        default=1,
        help="Number of successive training steps for the autoencoder."
    )

    options_parser.add_argument(
        "--num_discriminator_steps",
        type=int,
        default=1,
        help="Number of successive training steps for the discriminator."
    )

    options_parser.add_argument(
        "--num_generator_steps",
        type=int,
        default=1,
        help="Number of successive training steps for the generator."
    )

    options_parser.add_argument(
        "--autoencoder_noise_radius",
        type=float,
        default=0,
        help="Gaussian noise standard deviation for the latent code (autoencoder regularization)."
    )

    options_parser.add_argument(
        "--autoencoder_noise_anneal",
        type=float,
        default=0.995,
        help="Anneal the noise radius by this value after every epoch."
    )

    options_parser.add_argument(
        "--temperature",
        type=float,
        default=None,
        help="Gumbel-Softmax temperature."
    )

    options_parser.add_argument(
        "--penalty",
        type=float,
        default=0.1,
        help="WGAN-GP gradient penalty lambda."
    )

    options_parser.add_argument("--seed", type=int, help="Random number generator seed.", default=42)

    options = options_parser.parse_args()

    if options.seed is not None:
        np.random.seed(options.seed)
        torch.manual_seed(options.seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(options.seed)

    features = loaders[options.data_format](options.data)
    data = Dataset(features)
    train_data, val_data = data.split(1.0 - options.validation_proportion)

    if options.metadata is not None and options.temperature is not None:
        variable_sizes = load_variable_sizes_from_metadata(options.metadata)
        temperature = options.temperature
    else:
        variable_sizes = None
        temperature = None

    autoencoder = AutoEncoder(
        features.shape[1],
        code_size=options.code_size,
        encoder_hidden_sizes=parse_int_list(options.encoder_hidden_sizes),
        decoder_hidden_sizes=parse_int_list(options.decoder_hidden_sizes),
        variable_sizes=variable_sizes
    )

    load_or_initialize(autoencoder, options.input_autoencoder)

    generator = Generator(
        options.noise_size,
        options.code_size,
        hidden_sizes=parse_int_list(options.generator_hidden_sizes),
        bn_decay=options.bn_decay
    )

    load_or_initialize(generator, options.input_generator)

    discriminator = Discriminator(
        options.code_size,
        hidden_sizes=parse_int_list(options.discriminator_hidden_sizes),
        bn_decay=0,  # no batch normalization for the critic
        critic=True
    )

    load_or_initialize(discriminator, options.input_discriminator)

    train(
        autoencoder,
        generator,
        discriminator,
        train_data,
        val_data,
        options.output_autoencoder,
        options.output_generator,
        options.output_discriminator,
        options.output_loss,
        batch_size=options.batch_size,
        start_epoch=options.start_epoch,
        num_epochs=options.num_epochs,
        num_ae_steps=options.num_autoencoder_steps,
        num_disc_steps=options.num_discriminator_steps,
        num_gen_steps=options.num_generator_steps,
        noise_size=options.noise_size,
        l2_regularization=options.l2_regularization,
        learning_rate=options.learning_rate,
        ae_noise_radius=options.autoencoder_noise_radius,
        ae_noise_anneal=options.autoencoder_noise_anneal,
        variable_sizes=variable_sizes,
        temperature=temperature,
        penalty=options.penalty
    )


if __name__ == "__main__":
    main()