gan_trainer.py

# Copyright (c) 2016 Baidu, Inc. All Rights Reserved
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import random
import numpy
import cPickle
import sys,os
from PIL import Image

from paddle.trainer.config_parser import parse_config
from paddle.trainer.config_parser import logger
import py_paddle.swig_paddle as api
import matplotlib.pyplot as plt

def plot2DScatter(data, outputfile):
    '''
    Plot the data as a 2D scatter plot and save to outputfile
    data needs to be two dimensinoal
    '''
    x = data[:, 0]
    y = data[:, 1]
    logger.info("The mean vector is %s" % numpy.mean(data, 0))
    logger.info("The std vector is %s" % numpy.std(data, 0))

    heatmap, xedges, yedges = numpy.histogram2d(x, y, bins=50)
    extent = [xedges[0], xedges[-1], yedges[0], yedges[-1]]

    plt.clf()
    plt.scatter(x, y)
    plt.savefig(outputfile, bbox_inches='tight')

def CHECK_EQ(a, b):
    assert a == b, "a=%s, b=%s" % (a, b)

def copy_shared_parameters(src, dst):
    '''
    copy the parameters from src to dst
    :param src: the source of the parameters
    :type src: GradientMachine
    :param dst: the destination of the parameters
    :type dst: GradientMachine
    '''
    src_params = [src.getParameter(i)
               for i in xrange(src.getParameterSize())]
    src_params = dict([(p.getName(), p) for p in src_params])


    for i in xrange(dst.getParameterSize()):
        dst_param = dst.getParameter(i)
        src_param = src_params.get(dst_param.getName(), None)
        if src_param is None:
            continue
        src_value = src_param.getBuf(api.PARAMETER_VALUE)
        dst_value = dst_param.getBuf(api.PARAMETER_VALUE)
        CHECK_EQ(len(src_value), len(dst_value))
        dst_value.copyFrom(src_value)
        dst_param.setValueUpdated()
        
def print_parameters(src):
    src_params = [src.getParameter(i)
               for i in xrange(src.getParameterSize())]

    print "***************"
    for p in src_params:
        print "Name is %s" % p.getName()
        print "value is %s \n" % p.getBuf(api.PARAMETER_VALUE).copyToNumpyArray()

def load_mnist_data(imageFile):
    f = open(imageFile, "rb")
    f.read(16)

    # Define number of samples for train/test
    if "train" in imageFile:
        n = 60000
    else:
        n = 10000
    
    data = numpy.fromfile(f, 'ubyte', count=n*28*28).reshape((n, 28*28))
    data = data / 255.0 * 2.0 - 1.0

    f.close()
    return data

def load_cifar_data(cifar_path):
    batch_size = 10000
    data = numpy.zeros((5*batch_size, 32*32*3), dtype = "float32")
    for i in range(1, 6):
        file = cifar_path + "/data_batch_" + str(i)
        fo = open(file, 'rb')
        dict = cPickle.load(fo)
        fo.close()
        data[(i - 1)*batch_size:(i*batch_size), :] = dict["data"]
    
    data = data / 255.0 * 2.0 - 1.0
    return data

# synthesize 2-D uniform data
def load_uniform_data():
    data = numpy.random.rand(1000000, 2).astype('float32')
    return data

def merge(images, size):
    if images.shape[1] == 28*28:
        h, w, c = 28, 28, 1
    else:
        h, w, c = 32, 32, 3
    img = numpy.zeros((h * size[0], w * size[1], c))
    for idx in xrange(size[0] * size[1]):
        i = idx % size[1]
        j = idx // size[1]
        img[j*h:j*h+h, i*w:i*w+w, :] = \
          ((images[idx, :].reshape((h, w, c), order="F").transpose(1, 0, 2) + 1.0) / 2.0 * 255.0)
    return img.astype('uint8')

def save_images(images, path):
    merged_img = merge(images, [8, 8])
    if merged_img.shape[2] == 1:
        im = Image.fromarray(numpy.squeeze(merged_img)).convert('RGB')
    else:
        im = Image.fromarray(merged_img, mode="RGB")
    im.save(path)
    
def get_real_samples(batch_size, data_np):
    return data_np[numpy.random.choice(data_np.shape[0], batch_size, 
                                       replace=False),:]
    
def get_noise(batch_size, noise_dim):
    return numpy.random.normal(size=(batch_size, noise_dim)).astype('float32')

def get_fake_samples(generator_machine, batch_size, noise):
    gen_inputs = api.Arguments.createArguments(1)
    gen_inputs.setSlotValue(0, api.Matrix.createDenseFromNumpy(noise))
    gen_outputs = api.Arguments.createArguments(0)
    generator_machine.forward(gen_inputs, gen_outputs, api.PASS_TEST)
    fake_samples = gen_outputs.getSlotValue(0).copyToNumpyMat()
    return fake_samples

def get_training_loss(training_machine, inputs):
    outputs = api.Arguments.createArguments(0)
    training_machine.forward(inputs, outputs, api.PASS_TEST)
    loss = outputs.getSlotValue(0).copyToNumpyMat()
    return numpy.mean(loss)

def prepare_discriminator_data_batch_pos(batch_size, data_np):
    real_samples = get_real_samples(batch_size, data_np)
    labels = numpy.ones(batch_size, dtype='int32')
    inputs = api.Arguments.createArguments(2)
    inputs.setSlotValue(0, api.Matrix.createDenseFromNumpy(real_samples))
    inputs.setSlotIds(1, api.IVector.createVectorFromNumpy(labels))
    return inputs

def prepare_discriminator_data_batch_neg(generator_machine, batch_size, noise):
    fake_samples = get_fake_samples(generator_machine, batch_size, noise)
    labels = numpy.zeros(batch_size, dtype='int32')
    inputs = api.Arguments.createArguments(2)
    inputs.setSlotValue(0, api.Matrix.createDenseFromNumpy(fake_samples))
    inputs.setSlotIds(1, api.IVector.createVectorFromNumpy(labels))
    return inputs

def prepare_generator_data_batch(batch_size, noise):
    label = numpy.ones(batch_size, dtype='int32')
    inputs = api.Arguments.createArguments(2)
    inputs.setSlotValue(0, api.Matrix.createDenseFromNumpy(noise))
    inputs.setSlotIds(1, api.IVector.createVectorFromNumpy(label))
    return inputs


def find(iterable, cond):
    for item in iterable:
        if cond(item):
            return item
    return None


def get_layer_size(model_conf, layer_name):
    layer_conf = find(model_conf.layers, lambda x: x.name == layer_name)
    assert layer_conf is not None, "Cannot find '%s' layer" % layer_name
    return layer_conf.size


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("-d", "--data_source", help="mnist or cifar or uniform")
    parser.add_argument("--use_gpu", default="1", 
                        help="1 means use gpu for training")
    parser.add_argument("--gpu_id", default="0", 
                        help="the gpu_id parameter")
    args = parser.parse_args()
    data_source = args.data_source
    use_gpu = args.use_gpu
    assert data_source in ["mnist", "cifar", "uniform"]
    assert use_gpu in ["0", "1"]

    if not os.path.exists("./%s_samples/" % data_source):
        os.makedirs("./%s_samples/" % data_source)

    if not os.path.exists("./%s_params/" % data_source):
        os.makedirs("./%s_params/" % data_source)
        
    api.initPaddle('--use_gpu=' + use_gpu, '--dot_period=10', '--log_period=100', 
                   '--gpu_id=' + args.gpu_id, '--save_dir=' + "./%s_params/" % data_source)
    
    if data_source == "uniform":
        conf = "gan_conf.py"
        num_iter = 10000
    else:
        conf = "gan_conf_image.py"
        num_iter = 1000
        
    gen_conf = parse_config(conf, "mode=generator_training,data=" + data_source)
    dis_conf = parse_config(conf, "mode=discriminator_training,data=" + data_source)
    generator_conf = parse_config(conf, "mode=generator,data=" + data_source)
    batch_size = dis_conf.opt_config.batch_size
    noise_dim = get_layer_size(gen_conf.model_config, "noise")
    
    if data_source == "mnist":
        data_np = load_mnist_data("./data/mnist_data/train-images-idx3-ubyte")
    elif data_source == "cifar":
        data_np = load_cifar_data("./data/cifar-10-batches-py/")
    else:
        data_np = load_uniform_data()
    
    # this creates a gradient machine for discriminator
    dis_training_machine = api.GradientMachine.createFromConfigProto(
        dis_conf.model_config)
    # this create a gradient machine for generator    
    gen_training_machine = api.GradientMachine.createFromConfigProto(
        gen_conf.model_config)

    # generator_machine is used to generate data only, which is used for
    # training discriminator
    logger.info(str(generator_conf.model_config))
    generator_machine = api.GradientMachine.createFromConfigProto(
        generator_conf.model_config)
    
    dis_trainer = api.Trainer.create(
        dis_conf, dis_training_machine)

    gen_trainer = api.Trainer.create(
        gen_conf, gen_training_machine)
    
    dis_trainer.startTrain()
    gen_trainer.startTrain()
    
    # Sync parameters between networks (GradientMachine) at the beginning
    copy_shared_parameters(gen_training_machine, dis_training_machine)
    copy_shared_parameters(gen_training_machine, generator_machine)
    
    # constrain that either discriminator or generator can not be trained
    # consecutively more than MAX_strike times
    curr_train = "dis"
    curr_strike = 0
    MAX_strike = 5
     
    for train_pass in xrange(100):
        dis_trainer.startTrainPass()
        gen_trainer.startTrainPass()
        for i in xrange(num_iter):
            # Do forward pass in discriminator to get the dis_loss
            noise = get_noise(batch_size, noise_dim)
            data_batch_dis_pos = prepare_discriminator_data_batch_pos(
                batch_size, data_np)
            dis_loss_pos = get_training_loss(dis_training_machine, data_batch_dis_pos)
            
            data_batch_dis_neg = prepare_discriminator_data_batch_neg(
                generator_machine, batch_size, noise)
            dis_loss_neg = get_training_loss(dis_training_machine, data_batch_dis_neg)            
                         
            dis_loss = (dis_loss_pos + dis_loss_neg) / 2.0
            
            # Do forward pass in generator to get the gen_loss
            data_batch_gen = prepare_generator_data_batch(
                    batch_size, noise)
            gen_loss = get_training_loss(gen_training_machine, data_batch_gen)
             
            if i % 100 == 0:
                print "d_pos_loss is %s     d_neg_loss is %s" % (dis_loss_pos, dis_loss_neg) 
                print "d_loss is %s    g_loss is %s" % (dis_loss, gen_loss)
            
            # Decide which network to train based on the training history
            # And the relative size of the loss        
            if (not (curr_train == "dis" and curr_strike == MAX_strike)) and \
               ((curr_train == "gen" and curr_strike == MAX_strike) or dis_loss > gen_loss):
                if curr_train == "dis":
                    curr_strike += 1
                else:
                    curr_train = "dis"
                    curr_strike = 1                
                dis_trainer.trainOneDataBatch(batch_size, data_batch_dis_neg)
                dis_trainer.trainOneDataBatch(batch_size, data_batch_dis_pos)               
                copy_shared_parameters(dis_training_machine, gen_training_machine)
 
            else:
                if curr_train == "gen":
                    curr_strike += 1
                else:
                    curr_train = "gen"
                    curr_strike = 1
                gen_trainer.trainOneDataBatch(batch_size, data_batch_gen)
                # TODO: add API for paddle to allow true parameter sharing between different GradientMachines 
                # so that we do not need to copy shared parameters. 
                copy_shared_parameters(gen_training_machine, dis_training_machine)
                copy_shared_parameters(gen_training_machine, generator_machine)
 
        dis_trainer.finishTrainPass()
        gen_trainer.finishTrainPass()
        # At the end of each pass, save the generated samples/images
        fake_samples = get_fake_samples(generator_machine, batch_size, noise)
        if data_source == "uniform":
            plot2DScatter(fake_samples, "./%s_samples/train_pass%s.png" % (data_source, train_pass))
        else:
            save_images(fake_samples, "./%s_samples/train_pass%s.png" % (data_source, train_pass))
    dis_trainer.finishTrain()
    gen_trainer.finishTrain()

if __name__ == '__main__':
    main()