gan_trainer.py 12.2 KB
Newer Older
X
xuwei06 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17
# 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
W
wangyang59 已提交
18 19 20
import cPickle
import sys,os
from PIL import Image
X
xuwei06 已提交
21 22 23 24

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

def plot2DScatter(data, outputfile):
W
wangyang59 已提交
28 29 30 31
    '''
    Plot the data as a 2D scatter plot and save to outputfile
    data needs to be two dimensinoal
    '''
32 33 34 35 36 37 38 39 40 41 42
    x = data[:, 0]
    y = data[:, 1]
    print "The mean vector is %s" % numpy.mean(data, 0)
    print "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')
X
xuwei06 已提交
43 44 45 46 47

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

def copy_shared_parameters(src, dst):
W
wangyang59 已提交
48 49 50 51 52 53 54
    '''
    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
    '''
X
xuwei06 已提交
55 56 57 58
    src_params = [src.getParameter(i)
               for i in xrange(src.getParameterSize())]
    src_params = dict([(p.getName(), p) for p in src_params])

59

X
xuwei06 已提交
60 61 62 63 64 65 66 67 68 69
    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()
70 71 72 73
        
def print_parameters(src):
    src_params = [src.getParameter(i)
               for i in xrange(src.getParameterSize())]
X
xuwei06 已提交
74

75 76 77 78
    print "***************"
    for p in src_params:
        print "Name is %s" % p.getName()
        print "value is %s \n" % p.getBuf(api.PARAMETER_VALUE).copyToNumpyArray()
X
xuwei06 已提交
79

W
wangyang59 已提交
80 81 82 83 84 85 86 87 88 89
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
    
90 91
    data = numpy.fromfile(f, 'ubyte', count=n*28*28).reshape((n, 28*28))
    data = data / 255.0 * 2.0 - 1.0
W
wangyang59 已提交
92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126

    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')

127
def save_images(images, path):
W
wangyang59 已提交
128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144
    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))
X
xuwei06 已提交
145 146 147
    gen_outputs = api.Arguments.createArguments(0)
    generator_machine.forward(gen_inputs, gen_outputs, api.PASS_TEST)
    fake_samples = gen_outputs.getSlotValue(0).copyToNumpyMat()
148 149 150 151 152 153 154 155
    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)

W
wangyang59 已提交
156 157
def prepare_discriminator_data_batch_pos(batch_size, data_np):
    real_samples = get_real_samples(batch_size, data_np)
158 159
    labels = numpy.ones(batch_size, dtype='int32')
    inputs = api.Arguments.createArguments(2)
W
wangyang59 已提交
160 161
    inputs.setSlotValue(0, api.Matrix.createDenseFromNumpy(real_samples))
    inputs.setSlotIds(1, api.IVector.createVectorFromNumpy(labels))
162 163
    return inputs

W
wangyang59 已提交
164 165
def prepare_discriminator_data_batch_neg(generator_machine, batch_size, noise):
    fake_samples = get_fake_samples(generator_machine, batch_size, noise)
166 167
    labels = numpy.zeros(batch_size, dtype='int32')
    inputs = api.Arguments.createArguments(2)
W
wangyang59 已提交
168 169
    inputs.setSlotValue(0, api.Matrix.createDenseFromNumpy(fake_samples))
    inputs.setSlotIds(1, api.IVector.createVectorFromNumpy(labels))
170
    return inputs
X
xuwei06 已提交
171

W
wangyang59 已提交
172
def prepare_generator_data_batch(batch_size, noise):
X
xuwei06 已提交
173 174
    label = numpy.ones(batch_size, dtype='int32')
    inputs = api.Arguments.createArguments(2)
W
wangyang59 已提交
175 176
    inputs.setSlotValue(0, api.Matrix.createDenseFromNumpy(noise))
    inputs.setSlotIds(1, api.IVector.createVectorFromNumpy(label))
X
xuwei06 已提交
177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193
    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():
W
wangyang59 已提交
194 195 196 197 198 199 200 201 202 203 204
    parser = argparse.ArgumentParser()
    parser.add_argument("-d", "--dataSource", help="mnist or cifar or uniform")
    parser.add_argument("--useGpu", default="1", 
                        help="1 means use gpu for training")
    parser.add_argument("--gpuId", default="0", 
                        help="the gpu_id parameter")
    args = parser.parse_args()
    dataSource = args.dataSource
    useGpu = args.useGpu
    assert dataSource in ["mnist", "cifar", "uniform"]
    assert useGpu in ["0", "1"]
205 206 207 208 209 210 211

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

    if not os.path.exists("./%s_params/" % dataSource):
        os.makedirs("./%s_params/" % dataSource)
        
W
wangyang59 已提交
212
    api.initPaddle('--use_gpu=' + useGpu, '--dot_period=10', '--log_period=100', 
213
                   '--gpu_id=' + args.gpuId, '--save_dir=' + "./%s_params/" % dataSource)
W
wangyang59 已提交
214 215 216 217 218 219 220 221 222 223 224
    
    if dataSource == "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=" + dataSource)
    dis_conf = parse_config(conf, "mode=discriminator_training,data=" + dataSource)
    generator_conf = parse_config(conf, "mode=generator,data=" + dataSource)
X
xuwei06 已提交
225 226
    batch_size = dis_conf.opt_config.batch_size
    noise_dim = get_layer_size(gen_conf.model_config, "noise")
W
wangyang59 已提交
227 228 229 230 231 232 233 234
    
    if dataSource == "mnist":
        data_np = load_mnist_data("./data/mnist_data/train-images-idx3-ubyte")
    elif dataSource == "cifar":
        data_np = load_cifar_data("./data/cifar-10-batches-py/")
    else:
        data_np = load_uniform_data()
    
235
    # this creates a gradient machine for discriminator
X
xuwei06 已提交
236 237
    dis_training_machine = api.GradientMachine.createFromConfigProto(
        dis_conf.model_config)
W
wangyang59 已提交
238
    # this create a gradient machine for generator    
X
xuwei06 已提交
239 240 241 242
    gen_training_machine = api.GradientMachine.createFromConfigProto(
        gen_conf.model_config)

    # generator_machine is used to generate data only, which is used for
243
    # training discriminator
X
xuwei06 已提交
244 245 246
    logger.info(str(generator_conf.model_config))
    generator_machine = api.GradientMachine.createFromConfigProto(
        generator_conf.model_config)
W
wangyang59 已提交
247
    
X
xuwei06 已提交
248 249 250 251 252
    dis_trainer = api.Trainer.create(
        dis_conf, dis_training_machine)

    gen_trainer = api.Trainer.create(
        gen_conf, gen_training_machine)
W
wangyang59 已提交
253
    
X
xuwei06 已提交
254 255
    dis_trainer.startTrain()
    gen_trainer.startTrain()
W
wangyang59 已提交
256 257
    
    # Sync parameters between networks (GradientMachine) at the beginning
258 259
    copy_shared_parameters(gen_training_machine, dis_training_machine)
    copy_shared_parameters(gen_training_machine, generator_machine)
W
wangyang59 已提交
260 261 262
    
    # constrain that either discriminator or generator can not be trained
    # consecutively more than MAX_strike times
263 264 265
    curr_train = "dis"
    curr_strike = 0
    MAX_strike = 5
W
wangyang59 已提交
266
     
267
    for train_pass in xrange(100):
X
xuwei06 已提交
268 269
        dis_trainer.startTrainPass()
        gen_trainer.startTrainPass()
W
wangyang59 已提交
270 271 272
        for i in xrange(num_iter):
            # Do forward pass in discriminator to get the dis_loss
            noise = get_noise(batch_size, noise_dim)
273
            data_batch_dis_pos = prepare_discriminator_data_batch_pos(
W
wangyang59 已提交
274
                batch_size, data_np)
275 276 277
            dis_loss_pos = get_training_loss(dis_training_machine, data_batch_dis_pos)
            
            data_batch_dis_neg = prepare_discriminator_data_batch_neg(
W
wangyang59 已提交
278
                generator_machine, batch_size, noise)
279
            dis_loss_neg = get_training_loss(dis_training_machine, data_batch_dis_neg)            
W
wangyang59 已提交
280
                         
281 282
            dis_loss = (dis_loss_pos + dis_loss_neg) / 2.0
            
W
wangyang59 已提交
283
            # Do forward pass in generator to get the gen_loss
284
            data_batch_gen = prepare_generator_data_batch(
W
wangyang59 已提交
285
                    batch_size, noise)
286
            gen_loss = get_training_loss(gen_training_machine, data_batch_gen)
W
wangyang59 已提交
287 288 289
             
            if i % 100 == 0:
                print "d_pos_loss is %s     d_neg_loss is %s" % (dis_loss_pos, dis_loss_neg) 
290
                print "d_loss is %s    g_loss is %s" % (dis_loss, gen_loss)
W
wangyang59 已提交
291 292 293 294 295
            
            # 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):
296 297 298 299 300 301
                if curr_train == "dis":
                    curr_strike += 1
                else:
                    curr_train = "dis"
                    curr_strike = 1                
                dis_trainer.trainOneDataBatch(batch_size, data_batch_dis_neg)
W
wangyang59 已提交
302
                dis_trainer.trainOneDataBatch(batch_size, data_batch_dis_pos)               
303
                copy_shared_parameters(dis_training_machine, gen_training_machine)
W
wangyang59 已提交
304
 
305 306 307 308 309 310 311 312 313
            else:
                if curr_train == "gen":
                    curr_strike += 1
                else:
                    curr_train = "gen"
                    curr_strike = 1
                gen_trainer.trainOneDataBatch(batch_size, data_batch_gen)    
                copy_shared_parameters(gen_training_machine, dis_training_machine)
                copy_shared_parameters(gen_training_machine, generator_machine)
W
wangyang59 已提交
314
 
X
xuwei06 已提交
315 316
        dis_trainer.finishTrainPass()
        gen_trainer.finishTrainPass()
W
wangyang59 已提交
317 318 319 320 321
        # At the end of each pass, save the generated samples/images
        fake_samples = get_fake_samples(generator_machine, batch_size, noise)
        if dataSource == "uniform":
            plot2DScatter(fake_samples, "./%s_samples/train_pass%s.png" % (dataSource, train_pass))
        else:
322
            save_images(fake_samples, "./%s_samples/train_pass%s.png" % (dataSource, train_pass))
X
xuwei06 已提交
323 324 325 326 327
    dis_trainer.finishTrain()
    gen_trainer.finishTrain()

if __name__ == '__main__':
    main()