6-Gans-03_CGans条件Gans

import tensorflow as tf
import numpy as np
from types import MethodType, FunctionType
import argparse
import math
import os
import cv2
from tensorflow.examples.tutorials.mnist.input_data import read_data_sets

"""
使用手写数据集实现：多GPU的条件生成对抗网络
"""


class Config:
    """
    定义模型超参数的类
    """
    def __init__(self):
        self.batch_size = 32
        self.lr = 0.0002
        self.epoches = 50
        self.save_path = './models/mnist07_cgan1'
        self.eps = 1e-5
        self.gpus = get_gpus()
        self.logdir = './log/mnist07'
        self.z_size = 16
        # 手写数据集的路径。
        self.samples_path = './mnist'

        self.__create_path_with_no_exist(self.save_path)
        self.__create_path_with_no_exist(self.logdir)

    def to_dict(self):
        result = {}
        for name in dir(self):
            if name.startswith('__'):
                continue
            attr = getattr(self, name)
            if isinstance(attr, MethodType) or isinstance(attr, FunctionType):
                continue
            print(name, type(attr))
            result[name] = attr
        return result

    def from_cmd_line(self):
        """
        从命令行读取超参数。
        :return:
        """
        parser = argparse.ArgumentParser()
        attrs = self.to_dict()
        for name in attrs:
            attr = attrs[name]
            parser.add_argument('--' + name, type=type(attr), help='Default to %s' % attr, default=attr)

        a = parser.parse_args()

        for name in dir(a):
            if name in attrs:
                setattr(self, name, getattr(a, name))

    def __repr__(self):
        result = '{'
        attrs = self.to_dict()
        for name in attrs:
            result += '%s: %s; ' % (name, attrs[name])
        return result + '}'

    def __create_path_with_no_exist(self, path):
        if not os.path.exists(path):
            os.makedirs(path)
            print('成功创建文件夹:{}'.format(path))


class Tensors:
    """
    构建模型的主类。
    """
    def __init__(self, config: Config):
        self.config = config
        self.sub_tensors = []
        self.lr = tf.placeholder(tf.float32, [], name='lr')
        self.bn_train = tf.placeholder_with_default(True, shape=None, name='bn_train')
        optimizer = tf.train.AdamOptimizer(self.lr)

        with tf.variable_scope('Network'):
            for gpu_index in range(config.gpus):
                sub_tensors = SubTensors(self.config, optimizer, self.bn_train)
                self.sub_tensors.append(sub_tensors)
                # todo 模型重用
                tf.get_variable_scope().reuse_variables()

        self.train_opt = []
        self.loss = []
        # 在gpu 0 上面执行梯度下降
        with tf.device('/gpu:0'), tf.variable_scope('train'):
            for i in range(1, 4):
                # 获取子类中 名字为gradient%d的属性
                grads_and_vars = avg_gradients(self._get_sub_attrs('gradient%d' % i))
                self.train_opt.append(optimizer.apply_gradients(grads_and_vars))
                loss = tf.reduce_mean(self._get_sub_attrs('loss%d' % i))
                self.loss.append(loss)
                tf.summary.scalar('loss_%d' % i, loss)
            self.summary_op = tf.summary.merge_all()

    def _get_sub_attrs(self, name):
        """
        按照类中属性的名字来获取类的变量
        :param name:
        :return:
        """
        attrs = []
        for sub_tensor in self.sub_tensors:
            attr = getattr(sub_tensor, name)
            attrs.append(attr)
        return attrs


def avg_gradients(gradients_s):
    """
    多gpu求平均梯度值的方法
    :param gradients_s:
    :return:
    """
    rezult = {}
    for grad_and_var in gradients_s:
        for grad, var in grad_and_var:
            if var not in rezult:
                rezult[var] = []
            rezult[var].append(grad)
    avg_grads = [(tf.reduce_mean(rezult[var], axis=0), var) for var in rezult]
    return avg_grads

class SubTensors:
    """
    构建单个gpu上的模型。
    """
    def __init__(self, config: Config, optimizer, bn_train: bool):
        # 构建真实数据的占位符
        self.sample = tf.placeholder(tf.float32, [None, 28*28], name='sample')
        sample = tf.reshape(self.sample, shape=[-1, 28, 28, 1])

        self.z = tf.placeholder(tf.float32, [None, config.z_size], name='z')
        self.label = tf.placeholder(tf.int32, [None], name='label')
        label = tf.one_hot(self.label, depth=10)

        # 用生成网络生成假图片。
        with tf.variable_scope('generator'):
            fake_images = generator(self.z, label, alpha=0.2, bn_train=bn_train)
            self.fake_images = fake_images*255 # 用于可视化的。

        # 用判别网络进行鉴别
        with tf.variable_scope('discriminator'):
            sample_p = discriminator(sample, label, alpha=0.2, bn_train=bn_train)
            # todo 重用判别网络的模型参数，对假图片进行鉴别。
            tf.get_variable_scope().reuse_variables()
            fake_p = discriminator(fake_images, label, alpha=0.2, bn_train=bn_train)

        with tf.variable_scope('loss_gradient'):
            # 防止sample_p 为0
            sample_p = tf.maximum(sample_p, config.eps)
            self.loss1 = -tf.reduce_mean(tf.log(sample_p))

            fake_p = tf.maximum(fake_p, config.eps)
            self.loss2 = -tf.reduce_mean(1-tf.log(fake_p))
            self.loss3 = -tf.reduce_mean(tf.log(fake_p))

            # 分别计算梯度值
            vars_list = tf.trainable_variables()
            d_vars = [var for var in vars_list if 'discriminator' in var.name]
            g_vars = [var for var in vars_list if 'generator' in var.name]

            self.gradient1 = optimizer.compute_gradients(self.loss1, var_list=d_vars)
            self.gradient2 = optimizer.compute_gradients(self.loss2, var_list=d_vars)
            self.gradient3 = optimizer.compute_gradients(self.loss3, var_list=g_vars)


def discriminator(x, label, alpha=0.2, bn_train=True):
    """
    :param x:  [N, 28, 28, 1]
    :param label: [N, 10]
    :param alpha:
    :param bn_train:
    :return:
    """
    # 第一层不用BN
    x= tf.layers.conv2d(x, 32, 3, 1, 'same', name='conv1')
    x = tf.nn.leaky_relu(x, alpha=alpha)
    # [N, 28, 28, 32]

    filters = 32
    size = 28
    for i in range(2):
        filters *=2
        size //= 2
        x = tf.layers.conv2d(x, filters, 3, 2, 'same', use_bias=False, name='conv2_%d' % i)
        x = tf.layers.batch_normalization(x, training=bn_train, name='conv2_bn_%d' % i)
        x = tf.nn.leaky_relu(x, alpha=alpha)

        # 加入条件标签的信息
        lab = tf.layers.dense(label, size * size * filters, name='lab_dense1_%d' % i)
        lab = tf.reshape(lab, shape=[-1, size, size, filters])
        x += lab
    # [N, 14, 14, 64]
    # [N, 7, 7, 128]
    x = tf.layers.flatten(x)
    x = tf.layers.dense(x, 500, use_bias=False, name='dense1')
    x = tf.layers.batch_normalization(x, training=bn_train, name='dense1_bn')
    x = tf.nn.leaky_relu(x, alpha=alpha)
    # [N, 500]

    lab = tf.layers.dense(label, 500, name='lab_dense2')
    x += lab
    x = tf.layers.dense(x, 1, name='logits')
    x = tf.nn.sigmoid(x)
    return x


def generator(z, label, alpha=0.2, bn_train=True):
    """
    :param z:  [N, 16]
    :param label: [N, 10]
    :param alpha:
    :param bn_train:
    :return:
    """
    z = tf.layers.dense(z, 7*7*128, name='dense1')
    z = tf.reshape(z, shape=[-1, 7, 7, 128])
    z = tf.layers.batch_normalization(z, training=bn_train, name='dense1_bn')
    z = tf.nn.leaky_relu(z, alpha=alpha)
    # [N, 7, 7, 128]

    filters = 128
    size = 7
    for i in range(2):
        filters //= 2
        size *=2
        z = tf.layers.conv2d_transpose(
            z, filters, 3, 2, 'same', use_bias=False, name='deconv_%d' % i)
        z = tf.layers.batch_normalization(z, training=bn_train, name='deconv_bn_%d' % i)
        z = tf.nn.leaky_relu(z, alpha=alpha)

        # 加入条件标签的信息
        lab = tf.layers.dense(label, size*size*filters, name='dense2_%d' % i)
        lab = tf.reshape(lab, shape=[-1, size, size, filters])
        z += lab

    # [N, 14, 14, 64]
    # [N, 28, 28, 32]
    # 输出层  得到logits
    logits = tf.layers.conv2d_transpose(
            z, 1, 3, strides=1, padding='same', name='deconv_logits')
    # [N, 28, 28, 1]
    lab = tf.layers.dense(label, size**2, name='dense3')
    lab = tf.reshape(lab, shape=[-1, size, size, 1])
    logits += lab
    out = tf.nn.sigmoid(logits)
    return out


class Samples:
    """
    读取数据的类
    """
    def __init__(self, config: Config):
        self.config = config
        self.data = read_data_sets(self.config.samples_path)
        self.num = self.data.train.num_examples
    def next_batch(self, batch_size):
        return self.data.train.next_batch(batch_size)


class MnistCGAN:
    """
    执行会话的
    """
    def __init__(self):
        graph = tf.Graph()
        with graph.as_default():
            self.config = Config()
            # 从命令行读取参数
            self.config.from_cmd_line()
            self.tensors = Tensors(self.config)
            self.samples = Samples(self.config)

            cfg = tf.ConfigProto()
            cfg.allow_soft_placement = True
            self.sess = tf.Session(config=cfg, graph=graph)
            self.saver = tf.train.Saver(max_to_keep=1)

            # 恢复模型
            ckpt = tf.train.get_checkpoint_state(self.config.save_path)
            try:
                self.saver.restore(self.sess, ckpt.model_checkpoint_path)
                print('Load old model!')
            except:
                self.sess.run(tf.global_variables_initializer())
                print('No old model! train from scratch!')

    def train(self):
        file_writer = tf.summary.FileWriter(self.config.logdir, self.sess.graph)

        for e in range(1, self.config.epoches):
            steps_per_batch = self.samples.num // (self.config.gpus * self.config.batch_size)
            for step in range(steps_per_batch):
                feed = {self.tensors.lr: self.config.lr}

                # 循环喂入模型数据
                for gpu_index in range(self.config.gpus):
                    images, labels = self.samples.next_batch(self.config.batch_size)
                    feed[self.tensors.sub_tensors[gpu_index].sample] = images
                    feed[self.tensors.sub_tensors[gpu_index].label] = labels
                    feed[self.tensors.sub_tensors[gpu_index].z] = \
                        np.random.normal(size=[self.config.batch_size, self.config.z_size])
                losses = []
                for i in range(3):
                    _, loss = self.sess.run(
                        [self.tensors.train_opt[i], self.tensors.loss[i]], feed)
                    losses.append(loss)
                print('Epochs:{} - Step:{} - Loss1:{} - Loss2:{} - Loss3:{}'.format(
                    e, step, losses[0], losses[1], losses[2]
                ))

            if e % 5 ==0:
                save_files = os.path.join(self.config.save_path, 'model.ckpt')
                self.saver.save(self.sess, save_path=save_files)
                print('model saved to path:{}'.format(save_files))

    def predict(self):
        # 这里输入的标签，是循环从0--9依次生成的
        feed = {
            self.tensors.sub_tensors[0].z: np.random.normal(size=[self.config.batch_size, self.config.z_size]),
            self.tensors.sub_tensors[0].label: [e % 10 for e in range(self.config.batch_size)]
        }
        fake = self.sess.run(self.tensors.sub_tensors[0].fake_images, feed)
        fake = np.reshape(fake, newshape=[-1, 28])
        cv2.imwrite('output.jpg', fake)
        print('done')


def get_gpus():
    devices = os.getenv('CUDA_VISIBLE_DEVICES', '0')
    print(devices, flush=True)
    return len(devices.split(','))


if __name__ == '__main__':
    app = MnistCGAN()
    app.train()
    app.predict()