BigGAN

BigGAN

背景介绍

  BigGAN:于2019年发表于ICLR，被誉为史上最强GAN图像生成器，BigGAN为什么那么强，因为引入了大量的黑科技，目前GitHub上面都是基于Pytorch的实现，而且代码特别繁琐，这里向小伙伴们介绍我的TensorFlow2.0简易版实现。

BigGAN的特点

  **BigGAN可以认为SNGAN和SAGAN的结合，为了解决WGAN中的1-Lipshcitz问题，BigGAN在生成器和判别器中都借鉴了SNGAN中的Spectral Normalization(频谱归一化)的思想，而且借鉴了SAGAN的Self-Attention(注意力机制)**。
  Truncation Trick(截断技巧)，将噪声向量进行截断，可以提高样本的质量，但是降低了样本的多样性。
  Orthogonal Regularization(正交正则化)，可以降低权重系数之间的干扰。
  Class-Conditional-BatchNorm(类条件批归一化)，在归一化时引入分类信息，可以生成指定类型的图像。
  Hierarchical latent spaces(分层潜在空间)，输入噪声分布在网络的各个层，并不只作用于第一层。
  使用了ResNet网络结构，其中有输入和输出尺寸相同的Resblock层，输入尺寸宽高缩小两倍的Resblock_down层和输入尺寸宽高增大两倍的Resblock_up层。
  batch大，参数量大，训练时间长，在这里我只展示网络结构和一些细节，训练过程我就跳过了。

128x128网络结构

BigGAN图像分析

TensorFlow2.0实现

from functools import reduce
import tensorflow as tf
try:
    import tensorflow.python.keras as keras
except:
    import tensorflow.keras as keras


def compose(*funcs):
    if funcs:
        return reduce(lambda f, g: lambda *a, **kw: g(f(*a, **kw)), funcs)
    else:
        raise ValueError('Composition of empty sequence not supported.')


def orthogonal_regularizer(scale):

    def ortho_reg(w):
        shape = w.get_shape().as_list()
        c = shape[-1]

        w = tf.reshape(w, [-1, c])

        identity = tf.eye(c)

        w_transpose = tf.transpose(w)
        w_mul = tf.matmul(w_transpose, w)
        reg = tf.subtract(w_mul, identity)

        ortho_loss = tf.nn.l2_loss(reg)

        return scale * ortho_loss

    return ortho_reg


def orthogonal_regularizer_fully(scale):

    def ortho_reg_fully(w):
        _, c = w.get_shape().as_list()

        identity = tf.eye(c)
        w_transpose = tf.transpose(w)
        w_mul = tf.matmul(w_transpose, w)
        reg = tf.subtract(w_mul, identity)

        ortho_loss = tf.nn.l2_loss(reg)

        return scale * ortho_loss

    return ortho_reg_fully


class SpectralNorm(keras.layers.Layer):
    def __init__(self, iteration=1, **kwargs):
        super(SpectralNorm, self).__init__(**kwargs, dynamic=True)
        self.iteration = iteration

    def build(self, input_shape):
        self.u = self.add_variable(shape=[1, input_shape[-1]],
                                   initializer=tf.initializers.TruncatedNormal(1.),
                                   trainable=False)

    def call(self, inputs, **kwargs):
        shape = tf.shape(inputs)
        w = tf.reshape(inputs, shape=[-1, shape[-1]])
        u_hat = self.u
        for i in range(self.iteration):
            v_hat = tf.nn.l2_normalize(tf.matmul(u_hat, tf.transpose(w)))
            u_hat = tf.nn.l2_normalize(tf.matmul(v_hat, w))

        u_hat = tf.stop_gradient(u_hat)
        v_hat = tf.stop_gradient(v_hat)

        sigma = tf.matmul(tf.matmul(v_hat, w), tf.transpose(u_hat))
        with tf.control_dependencies([self.u.assign(u_hat)]):
            w_norm = w / sigma
            w_norm = tf.reshape(w_norm, inputs.get_shape())
        return w_norm

    def compute_output_shape(self, input_shape):

        return input_shape


class ClassConditionalBatchNorm(keras.layers.Layer):

    def __init__(self, name):
        super(ClassConditionalBatchNorm, self).__init__()
        self._name = name

    def build(self, input_shape):
        self.beta_dense = keras.layers.Dense(units=input_shape[0][-1])
        self.gamma_dense = keras.layers.Dense(units=input_shape[0][-1])

    def call(self, inputs, is_training=True):

        x, condition = inputs
        #
        split = keras.layers.Flatten()(condition)
        beta = self.beta_dense(split)
        gamma = self.gamma_dense(split)

        beta = tf.reshape(beta, shape=[-1, 1, 1, x.shape[-1]])
        gamma = tf.reshape(gamma, shape=[-1, 1, 1, x.shape[-1]])

        batch_mean, batch_var = tf.nn.moments(x, [0, 1, 2], keepdims=True)

        return (x - batch_mean) / batch_var * gamma + beta


class MyConv(keras.layers.Layer):
    def __init__(self, filters, kernel_size, strides, padding, name):
        super(MyConv, self).__init__()
        self._name = name
        self.filters = filters
        self.kernel_size = kernel_size
        self.strides = strides
        self.padding = padding

    def build(self, input_shape):
        self.w = self.add_weight(name='kernel',
                                 shape=(self.kernel_size, self.kernel_size, input_shape[-1], self.filters),
                                 initializer=weight_init, regularizer=weight_regularizer)
        self.b = self.add_weight(name='bias', shape=(self.filters,), initializer=keras.initializers.Zeros())

        if self._name.find('sn') != -1:
            self.u = self.add_weight(shape=[1, self.w.shape[-1]], initializer=tf.initializers.TruncatedNormal(1.))

    def call(self, inputs, **kwargs):
        if self._name.find('sn') != -1:
            shape = tf.shape(self.w)
            w = tf.reshape(self.w, shape=[-1, shape[-1]])
            u_hat = self.u
            for i in range(1):
                v_hat = tf.nn.l2_normalize(tf.matmul(u_hat, tf.transpose(w)))
                u_hat = tf.nn.l2_normalize(tf.matmul(v_hat, w))

            sigma = tf.matmul(tf.matmul(v_hat, w), tf.transpose(u_hat))

            with tf.control_dependencies([self.u.assign(u_hat)]):
                w_norm = w / sigma
                w_norm = tf.reshape(w_norm, self.w.get_shape())

            return tf.nn.bias_add(tf.nn.conv2d(inputs, w_norm, (1, self.strides, self.strides, 1), self.padding), self.b)

        return tf.nn.bias_add(tf.nn.conv2d(inputs, self.w, (1, self.strides, self.strides, 1), self.padding), self.b)


class MyDense(keras.layers.Layer):
    def __init__(self, units, name):
        super(MyDense, self).__init__()
        self._name = name
        self.units = units

    def build(self, input_shape):
        self.w = self.add_weight(name='kernel', shape=(input_shape[-1], self.units),
                                 initializer=weight_init, regularizer=weight_regularizer_fully)
        self.b = self.add_weight(name='bias', shape=(self.units,), initializer=keras.initializers.Zeros())

        if self._name.find('sn') != -1:
            self.u = self.add_weight(shape=[1, self.w.shape[-1]], initializer=tf.initializers.TruncatedNormal(1.))

    def call(self, inputs, **kwargs):
        if self._name.find('sn') != -1:
            shape = tf.shape(self.w)
            w = tf.reshape(self.w, shape=[-1, shape[-1]])
            u_hat = self.u
            for i in range(1):
                v_hat = tf.nn.l2_normalize(tf.matmul(u_hat, tf.transpose(w)))
                u_hat = tf.nn.l2_normalize(tf.matmul(v_hat, w))

            sigma = tf.matmul(tf.matmul(v_hat, w), tf.transpose(u_hat))

            with tf.control_dependencies([self.u.assign(u_hat)]):
                w_norm = w / sigma
                w_norm = tf.reshape(w_norm, self.w.get_shape())

            return tf.matmul(inputs, w_norm) + self.b

        return tf.matmul(inputs, self.w) + self.b


class Resblock(keras.layers.Layer):
    def __init__(self, filters, name):
        super(Resblock, self).__init__()
        self._name = name
        self.block = keras.Sequential([MyConv(filters, 3, 1, 'SAME', name='{}_part1_snconv1'.format(name)),
                                       keras.layers.BatchNormalization(momentum=0.8, name='{}_bn1'.format(name)),
                                       keras.layers.ReLU(name='{}_relu'.format(name)),
                                       MyConv(filters, 3, 1, 'SAME', name='{}_part1_snconv2'.format(name)),
                                       keras.layers.BatchNormalization(momentum=0.8, name='{}_bn2'.format(name))])
        self.add = keras.layers.Add(name='{}_add'.format(name))

    def call(self, inputs, **kwargs):
        x = self.block(inputs)
        output = self.add([x, inputs])

        return output


class Resblock_Down(keras.layers.Layer):
    def __init__(self, filters, name):
        super(Resblock_Down, self).__init__()
        self._name = name
        self.block1 = keras.Sequential([keras.layers.ReLU(name='{}_part1_relu1'.format(name)),
                                        MyConv(filters, 3, 1, 'SAME', name='{}_part1_snconv1'.format(name)),
                                        keras.layers.ReLU(name='{}_part1_relu2'.format(name)),
                                        MyConv(filters, 3, 1, 'SAME', name='{}_part1_snconv2'.format(name)),
                                        keras.layers.AveragePooling2D((2, 2), name='{}_part1_averagepool'.format(name))])

        self.block2 = keras.Sequential([MyConv(filters, 1, 1, 'SAME', name='{}_part2_snconv'.format(name)),
                                        keras.layers.AveragePooling2D((2, 2), name='{}_part2_averagepool'.format(name))])

        self.add = keras.layers.Add(name='{}_add'.format(name))

    def call(self, inputs, **kwargs):
        x1 = self.block1(inputs)
        x2 = self.block2(inputs)
        output = self.add([x1, x2])

        return output


class Resblock_Up(keras.layers.Layer):
    def __init__(self, filters, name):
        super(Resblock_Up, self).__init__()
        self._name = name
        self.cbn1 = ClassConditionalBatchNorm(name='{}_part1_cbn1'.format(name))
        self.cbn2 = ClassConditionalBatchNorm(name='{}_part1_cbn2'.format(name))
        self.block1_1 = keras.Sequential([keras.layers.ReLU(name='{}_part1_relu1'.format(name)),
                                          keras.layers.UpSampling2D((2, 2), name='{}_part1_upsampling'.format(name)),
                                          MyConv(filters, 3, 1, 'SAME', name='{}_part1_snconv1'.format(name))])

        self.block1_2 = keras.Sequential([keras.layers.ReLU(name='{}_part1_relu2'.format(name)),
                                          MyConv(filters, 3, 1, 'SAME', name='{}_part1_snconv2'.format(name))])

        self.block2 = keras.Sequential([keras.layers.UpSampling2D((2, 2), name='{}_part2_upsampling'.format(name)),
                                        MyConv(filters, 1, 1, 'SAME', name='{}_part2_snconv1'.format(name))])

        self.add = keras.layers.Add(name='{}_add'.format(name))

    def call(self, inputs, **kwargs):
        x, z = inputs
        x1 = self.cbn1([x, z])
        x1 = self.block1_1(x1)
        x1 = self.cbn2([x1, z])
        x1 = self.block1_2(x1)
        x2 = self.block2(x)
        output = self.add([x1, x2])

        return output


class SAblock(keras.layers.Layer):
    def __init__(self, filters, name):
        super(SAblock, self).__init__()
        self._name = name
        self.theta = MyConv(filters // 8, 1, 1, 'SAME', name='{}_theta'.format(name))
        self.phi = MyConv(filters // 8, 1, 1, 'SAME', name='{}_phi'.format(name))
        self.g = MyConv(filters, 1, 1, 'SAME', name='{}_g'.format(name))
        self.o = MyConv(filters, 1, 1, 'SAME', name='{}_conv4'.format(name))
        self.gamma = tf.Variable([0.])

    def call(self, inputs, **kwargs):
        theta = self.theta(inputs)
        theta = tf.reshape(theta, (-1, theta.shape[1] * theta.shape[2], theta.shape[-1]), name='{}_theta_reshape'.format(self._name))
        phi = self.phi(inputs)
        phi = tf.reshape(phi, (-1, phi.shape[1] * phi.shape[2], phi.shape[-1]), name='{}_phi_reshape'.format(self._name))
        g = self.g(inputs)
        g = tf.reshape(g, (-1, g.shape[1] * g.shape[2], g.shape[-1]), name='{}_g_reshape'.format(self._name))
        theta_phi = tf.matmul(theta, phi, transpose_b=True, name='{}_theta_dot_phi'.format(self._name))
        theta_phi = tf.nn.softmax(theta_phi, name='{}_softmax'.format(self.name))
        theta_phi_g = tf.matmul(theta_phi, g, name='{}_theta_phi_dot_g'.format(self._name))
        theta_phi_g = tf.reshape(theta_phi_g, shape=(-1, inputs.shape[1], inputs.shape[2], inputs.shape[3]), name='{}_theta_phi_g_reshape'.format(self._name))
        o = self.o(theta_phi_g)

        return o * self.gamma + inputs


def generator(input_shape_noise, input_shape_label):
    input_tensor_noise = keras.layers.Input(input_shape_noise, name='input_noise')
    input_tensor_label = keras.layers.Input(input_shape_label, name='input_label')

    embedding_tensor = compose(keras.layers.Embedding(1000, 120, name='embedding'),
                               keras.layers.Flatten(name='flatten'))(input_tensor_label)

    noise_split = tf.split(input_tensor_noise, 6, -1, name='split')
    for i in range(1, 6):
        noise_split[i] = keras.layers.Concatenate(name='concatenate{}'.format(i + 1))([noise_split[i], embedding_tensor])

    x = compose(keras.layers.Dense(1024 * 16, activation='relu', name='dense_relu'),
                keras.layers.Reshape((4, 4, 1024), name='reshape'))(noise_split[0])
    x = Resblock_Up(1024, name='resblockup1')([x, noise_split[1]])
    x = Resblock_Up(512, name='resblockup2')([x, noise_split[2]])
    x = Resblock_Up(256, name='resblockup3')([x, noise_split[3]])
    x = Resblock_Up(128, name='resblockup4')([x, noise_split[4]])
    x = SAblock(128, name='sablock')(x)
    x = Resblock_Up(64, name='resblockup5')([x, noise_split[5]])
    x = compose(keras.layers.BatchNormalization(momentum=0.8, name='bn'),
                keras.layers.ReLU(name='relu'),
                MyConv(3, 3, 1, 'SAME', name='conv'),
                keras.layers.Activation('tanh', name='tanh'))(x)

    model = keras.Model([input_tensor_noise, input_tensor_label], x, name='BigGAN-Generator')

    return model


def discriminator(input_shape_image, input_shape_label):
    input_tensor_image = keras.layers.Input(input_shape_image, name='input_image')
    input_tensor_label = keras.layers.Input(input_shape_label, name='input_label')

    x = compose(Resblock_Down(64, name='resblockdown1'),
                SAblock(64, name='sablock'),
                Resblock_Down(128, name='resblockdown2'),
                Resblock_Down(256, name='resblockdown3'),
                Resblock_Down(512, name='resblockdown4'),
                Resblock_Down(1024, name='resblockdown5'),
                Resblock(1024, name='resblock6'))(input_tensor_image)

    x = tf.reduce_sum(x, axis=[1, 2], name='global_sumpool')
    output_tensor = keras.layers.Dense(1, name='dense')(x)

    embedding_tensor = compose(keras.layers.Embedding(1000, 1024, name='embedding'),
                               keras.layers.Flatten(name='flatten'))(input_tensor_label)
    output_tensor = output_tensor + tf.reduce_sum(embedding_tensor * x, 1, keepdims=True, name='reduce_sum')

    model = keras.Model([input_tensor_image, input_tensor_label], output_tensor, name='BigGAN-Discriminator')

    return model


def biggan(input_shape_noise, input_shape_image, input_shape_label, model_g, model_d):
    input_noise = keras.layers.Input(input_shape_noise, name='input_noise')
    input_real_image = keras.layers.Input(input_shape_image, name='input_image')
    input_label = keras.layers.Input(input_shape_label, name='input_label')

    model_g.trainable = False
    fake = model_g([input_noise, input_label])
    real_conf = model_d([input_real_image, input_label])
    fake_conf = model_d([fake, input_label])

    model_discriminator = keras.Model([input_noise, input_real_image, input_label], [real_conf, fake_conf], name='BigGAN-discriminator')
    model_discriminator.compile(optimizer=optimizer_d, loss=[d_loss, d_loss], loss_weights=[1, 1])

    model_g.trainable = True
    model_d.trainable = False

    model_generator = keras.Model([input_noise, input_label], fake_conf, name='BigGAN-generator')
    model_generator.compile(optimizer=optimizer_g, loss=g_loss)

    return model_generator, model_discriminator


def d_loss(y_true, y_pred):

    return tf.reduce_mean(tf.nn.relu(1 - y_true * y_pred))


def g_loss(y_true, y_pred):

    return -tf.reduce_mean(y_pred)


if __name__ == '__main__':
    weight_init = tf.initializers.TruncatedNormal(mean=0.0, stddev=0.02)
    weight_regularizer = orthogonal_regularizer(0.0001)
    weight_regularizer_fully = orthogonal_regularizer_fully(0.0001)

    optimizer_g = keras.optimizers.Adam(0.00005, 0, 0.999, epsilon=1e-5)
    optimizer_d = keras.optimizers.Adam(0.0002, 0, 0.999, epsilon=1e-5)

    model_d = discriminator(input_shape_image=(128, 128, 3), input_shape_label=(1,))

    model_g = generator(input_shape_noise=(120,), input_shape_label=(1,))

    model_g.build(input_shape=[(120,), (1,)])
    model_g.summary()
    keras.utils.plot_model(model_g, 'BigGAN-generator.png', show_shapes=True, show_layer_names=True)

    model_d.build(input_shape=[(128, 128, 3), (1,)])
    model_d.summary()
    keras.utils.plot_model(model_d, 'BigGAN-discriminator.png', show_shapes=True, show_layer_names=True)

    model_generator, model_discriminator = biggan(input_shape_noise=(120,), input_shape_image=(128, 128, 3), input_shape_label=(1,), model_g=model_g, model_d=model_d)

    model_generator.build(input_shape=[(120,), (1,)])
    model_generator.summary()
    keras.utils.plot_model(model_generator, 'BigGAN-generate.png', show_shapes=True, show_layer_names=True)

    model_discriminator.build(input_shape=[(120,), (128, 128, 3), (1,)])
    model_discriminator.summary()
    keras.utils.plot_model(model_discriminator, 'BigGAN-discriminate.png', show_shapes=True, show_layer_names=True)

模型运行结果

小技巧

1. 图像输入可以先将其归一化到0-1之间或者-1-1之间，因为网络的参数一般都比较小，所以归一化后计算方便，收敛较快。
2. 注意其中的一些维度变换和numpy，tensorflow常用操作，否则在阅读代码时可能会产生一些困难。
3. 可以设置一些权重的保存方式，学习率的下降方式和早停方式。
4. BigGAN对于网络结构，优化器参数，网络层的一些超参数都是非常敏感的，效果不好不容易发现原因，这可能需要较多的工程实践经验。
5. 先创建判别器，然后进行compile，这样判别器就固定了，然后创建生成器时，不要训练判别器，需要将判别器的trainable改成False，此时不会影响之前固定的判别器，这个可以通过模型的_collection_collected_trainable_weights属性查看，如果该属性为空，则模型不训练，否则模型可以训练，compile之后，该属性固定，无论后面如何修改trainable，只要不重新compile，都不影响训练。
6. 代码中正交正则化使用了闭包的概念，有关闭包的使用，可以参考我的另一篇博客，Closure & Decorators(闭包和装饰器)
7. 这个模型效果太好，生成的图片甚至比真实图片还要好，一些纹理，背景细节都可以完美呈现，但是想自己实现训练过程，非常困难，因此建议小伙伴了解就可以，不用亲自实践。

BigGAN小结

  BigGAN分为很多版本，有128的图像版本，256的图像版本和512的图像版本，具体模型结构都很类似，但是参数量指数级增长。这是最小的BigGAN版本，参数量都可以达到80M，虽然VGG16的参数量有一亿多，但是网络结构简单，因此训练反而快，而BigGAN含有很多细节操作，会花费较长的时间，因此训练起来非常慢，BigGAN不是单打独斗，在特点中已经分析了，可以看成SNGAN和SAGAN的共同作品，因此关于其中的数学推导可以参考网络的其他资源，在这里也不过多赘述，作为史上最强的GAN图像生成器，小伙伴们一定要了解它。