DCGAN实现

DCGAN实现

代码

• dcgan.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import math
import argparse
import cv2
import numpy as np
import tensorflow as tf

# DataManager负责提供数据
class DataManager(object):

    def __init__(self, data_dir):
        self.data_dir = data_dir
        self.im_shape = (48, 48, 3)
        self.im_list = self._get_im_names()
        self.batch_size = 64
        self.chunk_size = len(self.im_list) // self.batch_size

    def _get_im_names(self):
        if not self.data_dir:
            return np.asarray([])
        im_list = np.asarray(os.listdir(self.data_dir))
        np.random.shuffle(im_list)
        return im_list

    def imread(self, im_name):
        im = cv2.imread(os.path.join(self.data_dir, im_name))
        im = cv2.resize(im, self.im_shape[:2])
        im = (im.astype('float32') - 127.5) / 128.0
        return im

    def imwrite(self, name, im):
        im = (im * 128.0 + 127.5)
        im = im.astype('uint8')
        cv2.imwrite('./images/%s.jpg' % name, im)

    def next_batch(self):
        start = 0
        end = start + self.batch_size
        for i in range(self.chunk_size):
            name_list = self.im_list[start: end]
            batch_im_list = np.asarray([self.imread(im_name) for im_name in name_list])
            yield batch_im_list
            start += self.batch_size
            end += self.batch_size

# 不使用任何其他框架(Keras, Slim), 神经网络中所有的操作都重新封装成一个方法
class DCGAN(object):

    def __init__(self, data_dir):
        # 通过data_manager控制数据的输入与输出
        self.data_manager = DataManager(data_dir)
        self.batch_size = self.data_manager.batch_size
        self.im_shape = self.data_manager.im_shape
        self.chunk_size = self.data_manager.chunk_size

        # 噪声的长度
        self.z_len = 100
        self.learning_rate = 0.0002
        self.epochs = 100
        self.beta1 = 0.5
        self.sample_size = 64

    # 全连接层
    def fc(self, ims, output_size, scope='fc'):
        with tf.variable_scope(scope, reuse=False):
            weights = tf.get_variable('weights', [ims.shape[1], output_size], tf.float32, initializer=tf.truncated_normal_initializer(stddev=0.02))
            biases = tf.get_variable('biases', [1, output_size], initializer=tf.constant_initializer(0.0))
            return tf.matmul(ims, weights) + biases

    # 批量均值化
    def batch_norm(self, x, epsilon=1e-5, momentum=0.9, scope='batch_norm', is_training=True):
        with tf.variable_scope(scope, reuse=False):
            return tf.contrib.layers.batch_norm(x, epsilon=epsilon, decay=momentum, updates_collections=None, scale=True, is_training=is_training)

    # 卷积层
    def conv2d(self, ims, output_dim, scope='conv2d'):
        with tf.variable_scope(scope, reuse=False):
            # 在Tensorflow中, SAME不是一般人理解的SAME, 在此框架中, 只要知道了输入的维度和stride的大小, 让输入的维度除以stride的大小就是卷积之后的维度
            # 在卷积中, ksize的维度为[height, width, in_channels, out_channels], 注意: 与转置卷积不同
            ksize = [5, 5, ims.shape[-1], output_dim]
            strides = [1, 2, 2, 1]
            weights = tf.get_variable('weights', ksize, tf.float32, initializer=tf.truncated_normal_initializer(stddev=0.02))
            biases = tf.get_variable('biases', [1, 1, 1, output_dim], tf.float32, initializer=tf.constant_initializer(0.0))
            conv = tf.nn.conv2d(ims, weights, strides=strides, padding='SAME') + biases
            return conv

    # 转置卷积层
    def deconv2d(self, ims, output_shape, scope='deconv2d'):
        with tf.variable_scope(scope, reuse=False):
            ksize = [5, 5, output_shape[-1], ims.shape[-1]]
            strides = [1, 2, 2, 1]

            weights = tf.get_variable('weights', ksize, tf.float32, initializer=tf.truncated_normal_initializer(stddev=0.02))
            biases = tf.get_variable('biases', [1, 1, 1, output_shape[-1]], tf.float32, initializer=tf.constant_initializer(0.0))

            deconv = tf.nn.conv2d_transpose(ims, weights, output_shape=output_shape, strides=strides) + biases
            return deconv

    # leaky ReLu
    def lrelu(self, x, alpha=0.2):
        return tf.maximum(x, x * alpha)

    # 判别器, 比较简单, 就是传统的分类, 不过去掉了池化层, 添加了batch norm
    def discriminator(self, ims, reuse=False):
        with tf.variable_scope('discriminator', reuse=reuse):
            net = self.conv2d(ims, 64, scope='d_conv_1')
            net = self.lrelu(net) 

            net = self.conv2d(net, 64 * 2, scope='d_conv_2')
            net = self.batch_norm(net, scope='d_bn_2')
            net = self.lrelu(net)

            net = self.conv2d(net, 64 * 4, scope='d_conv_3')
            net = self.batch_norm(net, scope='d_bn_3')
            net = self.lrelu(net)

            net = self.conv2d(net, 64 * 8, scope='d_conv_4')
            net = self.batch_norm(net, scope='d_bn_4')
            net = self.lrelu(net)

            net = self.fc(tf.reshape(net, [-1, net.shape[1] * net.shape[2] * net.shape[3]]), 1, scope='d_fc_5')
            return tf.nn.sigmoid(net), net

    # 生成器, 就是一个解码器, 去掉了池化层, 添加了Bath norm, 左右的结果通过tanh输出
    def generator(self, noise_z, is_training=True):
        with tf.variable_scope('generator', reuse=False):
            # 训练输入的图像为48x48, 反过来计算出各个网络层的图像维度
            net = self.fc(noise_z, 3 * 3 * 64 * 8)
            net = tf.reshape(net, [-1, 3, 3, 64 * 8])
            net = self.batch_norm(net, scope='g_bn_1', is_training=is_training)
            net = tf.nn.relu(net)

            net = self.deconv2d(net, [self.batch_size, 6, 6, 64 * 4], scope='g_conv_2')
            net = self.batch_norm(net, scope='g_bn_2', is_training=is_training)
            net = tf.nn.relu(net)

            net = self.deconv2d(net, [self.batch_size, 12, 12, 64 * 2], scope='g_conv_3')
            net = self.batch_norm(net, scope='g_bn_3', is_training=is_training)
            net = tf.nn.relu(net)

            net = self.deconv2d(net, [self.batch_size, 24, 24, 64], scope='g_conv_4')
            net = self.batch_norm(net, scope='g_bn_4', is_training=is_training)
            net = tf.nn.relu(net)

            net = self.deconv2d(net, [self.batch_size, self.im_shape[0], self.im_shape[1], 3], scope='g_conv_5')

            return tf.nn.tanh(net)

    def train(self):
        real_ims = tf.placeholder(tf.float32, [self.batch_size, self.im_shape[0], self.im_shape[1], self.im_shape[2]], name='real_ims')
        noise_z = tf.placeholder(tf.float32, [None, self.z_len], name='noise_z')

        # Loss functions
        fake_ims = self.generator(noise_z)
        real_prob, real_logits = self.discriminator(real_ims)
        fake_prob, fake_logits = self.discriminator(fake_ims, reuse=True)

        real_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=real_logits, labels=tf.ones_like(real_logits)))
        fake_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=fake_logits, labels=tf.zeros_like(fake_logits)))
        g_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=fake_logits, labels=tf.ones_like(fake_logits)))
        d_loss = real_loss + fake_loss

        real_loss_sum = tf.summary.scalar('real_loss', real_loss)
        fake_loss_sum = tf.summary.scalar('fake_loss', fake_loss)
        g_loss_sum = tf.summary.scalar('g_loss', g_loss)
        d_loss_sum = tf.summary.scalar('d_loss', d_loss)

        # Optimizer
        train_vars = tf.trainable_variables()
        d_vars = [var for var in train_vars if var.name.startswith('discriminator')]
        g_vars = [var for var in train_vars if var.name.startswith('generator')]

        d_global_step = tf.Variable(0, name='d_global_step', trainable=False)
        g_global_step = tf.Variable(0, name='d_global_step', trainable=False)
        d_optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate, beta1=self.beta1).minimize(d_loss, var_list=d_vars, global_step=d_global_step)
        g_optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate, beta1=self.beta1).minimize(g_loss, var_list=g_vars, global_step=g_global_step)

        saver = tf.train.Saver()
        init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())
        with tf.Session() as sess:
            sess.run(init_op)
            d_merged = tf.summary.merge([d_loss_sum, real_loss_sum, fake_loss_sum])
            g_merged = tf.summary.merge([g_loss_sum])
            writer = tf.summary.FileWriter('./logs', sess.graph)
            ckpt = tf.train.get_checkpoint_state('./checkpoints')
            if ckpt and ckpt.model_checkpoint_path:
                saver.restore(sess, ckpt.model_checkpoint_path)
                print('Restore from model.ckpt')
            else:
                print('Checkpoint is not found!')
            for epoch in range(self.epochs):
                batches = self.data_manager.next_batch()
                for batch in batches:
                    noises = np.random.uniform(-1, 1, size=(self.batch_size, self.z_len)).astype(np.float32)
                    _, d_summary, d_step = sess.run([d_optimizer, d_merged, d_global_step], feed_dict={real_ims: batch, noise_z: noises})
                    sess.run(g_optimizer, feed_dict={noise_z: noises})
                    _, g_summary, g_step = sess.run([g_optimizer, g_merged, g_global_step], feed_dict={noise_z: noises})

                    writer.add_summary(d_summary, d_step)
                    writer.add_summary(g_summary, g_step)

                    loss_d, loss_real, loss_fake, loss_g = sess.run([d_loss, real_loss, fake_loss, g_loss], feed_dict={real_ims: batch, noise_z: noises})

                    print('Epoch: %s, Dis Step: %s, d_loss: %s, real_loss: %s, fake_loss: %s, Gen Step: %s, g_loss: %s' \
                            % (epoch, d_step, loss_d, loss_real, loss_fake, g_step, loss_g))
                    if g_step % 100 == 0:
                        saver.save(sess, './checkpoints/model.ckpt', global_step=g_step)
                        print('G Step %s Save model' % g_step)

    def gen(self):
        noise_z = tf.placeholder(tf.float32, [None, self.z_len], name='noise_z')
        sample_ims = self.generator(noise_z, is_training=False)
        saver = tf.train.Saver()
        with tf.Session() as sess:
            sess.run(tf.global_variables_initializer())
            sess.run(tf.local_variables_initializer())
            sample_noise = np.random.uniform(-1, 1, size=(self.sample_size, self.z_len))
            ckpt = tf.train.get_checkpoint_state('./checkpoints')
            if ckpt and ckpt.model_checkpoint_path:
                saver.restore(sess, ckpt.model_checkpoint_path)
                samples = sess.run(sample_ims, feed_dict={noise_z: sample_noise})
                for idx, sample in enumerate(samples):
                    self.data_manager.imwrite(idx, sample)
            else:
                print('Checkpoint is not found!')
                return

def data_load_test():
    manager = DataManager()
    batch = manager.next_batch()
    im_list = next(batch)
    for idx, im in enumerate(im_list):
        manager.imwrite(idx, im)

def main(argv=None):
    parser = argparse.ArgumentParser()
    parser.add_argument('--train', help='path to dataset')
    parser.add_argument('--gen', help='path to store images')
    args = parser.parse_args()
    if args.train:
        dcgan = DCGAN(args.train)
        dcgan.train()
    elif args.gen:
        if args.gen == 'yes':
            dcgan = DCGAN(None)
            dcgan.gen()
        else:
            print('should be --gen yes')
    else:
        print('...')

if __name__ == '__main__':
    main()