【GANs】Wasserstein GAN

4 W-GAN

在生成对抗网络中，$JS$散度不适合衡量生成数据分布和真实数据分布的距离。由于通过优化交叉熵（$JS$散度）训练生成对抗网络会导致训练稳定性和模型坍塌问题，因此改进GAN，就需要改变其损失函数。

4.1 W-GAN简介

Wasserstein GAN原文链接

W-GAN通过使用$Wasserstein$距离代替优化$JS$散度来优化训练的生成对抗网络。
对于真实分布$p_r$和模型分布$p_{\theta }$，他们的$1st-Wasserstein$距离为：
$$\begin{array}{r}{\mathbf{W}}^1(p_r,p_{\theta })=\underset{\gamma \sim \Gamma (p_r,p_{\theta })}{\inf }{\mathrm{E}}_{(x,y)\sim \gamma }\left[\Vert x-y\Vert \right]\end{array}$$
其中$\Gamma (p_r,p_{\theta })$是边界分布为$p_r$和$p_{\theta }$的所有可能的联合分布集合。

当两个分布没有重叠或者重叠非常少时，他们之间的$KL$散度为$+\infty$，$JS$散度为$log2$，并不随着两个分布之间的距离而变化。而$1st-Wasserstein$距离依然可以衡量两个没有重叠分布之间的距离。

两个分布$p_r$和$p_{\theta }$的$1st-Wasserstein$距离通常难以直接计算，但是两个分布的$1st-Wasserstein$距离有一个对偶形式：
$$\begin{array}{r}{\mathbf{W}}^1(p_r,p_{\theta })=\underset{{\Vert f\Vert }_L\le 1}{\sup }\left({\mathrm{E}}_{x\sim p_r}\left[f(x)\right]-{\mathrm{E}}_{x\sim p_{\theta }}\left[f(x)\right]\right)\end{array}$$
其中$f:{\mathbb{R}}^d\to \mathbb{R}$为$1-Lipschitz$函数,满足：
$$\begin{array}{r}{\Vert f\Vert }_L\le 1\triangleq \underset{x\ne y}{\sup }\frac{\left|f(x)-f(y)\right|}{\left|x-y\right|}⩽1\end{array}$$
公式$(3)$称为$Kantorovich-Rubinstein$对偶定理。

根据$Kantorovich-Rubinstein$对偶定理，两个分布$p_r$和$p_{\theta }$之间的$1st-Wasserstein$距离可以转化为一个满足$1-Lipschitz$连续的函数在分布$p_r$和$p_{\theta }$下期望的差的上界。通常情况下，$1-Lipschitz$连续的约束可以宽松为$K-Lipschitz$连续。这样分布$p_r$和$p_{\theta }$之间的$1st-Wasserstein$距离为：
$$\begin{array}{r}{\mathbf{W}}^1(p_r,p_{\theta })=\frac{1}{K}\underset{{\Vert f\Vert }_L\le K}{\sup }\left({\mathrm{E}}_{x\sim p_r}\left[f(x)\right]-{\mathrm{E}}_{x\sim p_{\theta }}\left[f(x)\right]\right)\end{array}$$

评价网络

计算$(4)$的上界并不容易。根据神经网络的通用近似定理，假设存在一个神经网络使得可以达到这个上界。

令$f(x;\varphi )$为一个神经网络，假设存在参数集合$\Phi$，对于所有的$\varphi \in \Phi$，$f(x;\varphi )$为$K-Lipschitz$连续函数，那么$(4)$中的上界可以近似转换为：
$$\begin{array}{r}\underset{\varphi \in \Phi }{\max }\left({\mathrm{E}}_{x\sim p_r}\left[f(x;\varphi )\right]-{\mathrm{E}}_{x\sim p_{\theta }}\left[f(x;\varphi )\right]\right)\end{array}$$
其中$f(x;\varphi )$称为评价网络（Critic Network）。

与标准GAN中的判别网络的值域为$[0,1]$不同，评价网络$f(x;\varphi )$的最后一层为线性层，其值域没有限制。这样只需要找到一个网络$f(x;\varphi )$使其在两个分布$p_r$和$p_{\theta }$下的期望最大。即对于真实样本，$f(x;\varphi )$的打分要尽可能低。

为了使得$f(x;\varphi )$满足$K-Lipschitz$连续，一种近似的方法使限制参数的取值范围。因为神经网络为连续可导函数，满足$K-Lipschitz$连续可以近似为其关于$x$的偏导数的模$\Vert \frac{\partial f(x;\varphi )}{\partial x}\Vert$小于某个上界。由于这个偏导数的大小一般和参数的取值范围相关，我们可以通过限制参数$\varphi$的取值范围来近似，令$\varphi \in [-c,c]$，$c$为一个较小的正数，比如$0.01$。

生成网络

生成网络的目标是使得评价网络$f(x;\varphi )$对其生成样本的打分尽可能高，即
$$\begin{array}{r}\underset{\theta }{\max }{\mathrm{E}}_{z\sim p_{(z)}}\left[f(G(z;\theta );\varphi )\right]\end{array}$$
因为$f(x;\varphi )$为不饱和函数，所以生成网络参数$\theta$的梯度不会消失，理论上解决了原始GAN训练不稳定的问题。并且W-GAN中生成网络的目标函数不再是两个分布的比率，在一定程度上缓解了模型坍塌问题，使得生成的样本具有多样性。

4.2 散度

KL散度

$KL散度(Kullback-LeiblerDivergence)$也叫$KL距离$或$相对熵(RelativeEntropy)$，是用概率分布$q$来近似分布$p$时造成的信息的损失量。$KL散度$是按照概率分布$q$的最优编码来对真实分布为$p$的信息进行编码，其平均编码长度(即交叉熵)$H(p,q)$和$p$的最优平均编码长度(即熵)$H(p)$之间的差异。对于离散概率分布$p$和$q$，从$q$到$p$的$KL散度$定义为
$$\begin{array}{rl}KL(p,q)& =H(p,q)-H(p)\\ & =\sum _{x}p(x)\log \frac{p(x)}{q(x)}\end{array}$$
其中，为了保证连续性，定义$0log\frac{0}{0}=0,0log\frac{0}{q}=0$。

$KL散度$总是非负的，$KL(p,q)\ge 0$，可以衡量两个概率分布之间的距离。$KL散度$只有当$p=q$时，$KL(p,q)=0$。

如果两个分布越接近，$KL散度$越小；反之亦反。

但$KL散度$并不是一个真正的度量或距离，一是$KL散度$不满足距离的对称性，二是$KL散度$不满足距离的三角不等式性质。

JS散度

$JS散度(Jensen-ShannonDivergence)$是一种对称的衡量两个分布相似度的度量方式，定义为
$$\begin{array}{r}JS(p,q)=\frac{1}{2}KL(p,m)+\frac{1}{2}KL(q,m)\end{array}$$
其中$m=\frac{1}{2}(p+q)$

$JS散度$是$KL散度$一种改进，但是两种散度都存在一个问题，即如果两个分布$p,q$没有重叠或者重叠非常少，$KL散度$和$JS散度$都很难衡量两个分布的距离。

Wasserstein距离

$Wasserstein距离(WassersteinDistance)$也用于衡量两个分布之间的距离。对于两个分布$q_1,q_2$，$p^{th}-Wasserstein$距离定义为
$$\begin{array}{r}{\mathbf{W}}_p(q_1,q_2)=(\underset{\gamma (x,y)\in \Gamma (q_1,q_2)}{\inf }{\mathrm{E}}_{(x,y)\sim \gamma (x,y)}\left[d{(x,y)}^p\right]{)}^{\frac{1}{p}}\end{array}$$
其中$\Gamma (q_1,q_2)$是边际分布为$q_1$和$q_2$的所有可能的联合分布集合，$d(x,y)$为$x$和$y$的距离。

$Wasserstein距离$相比$JS散度$和$KL散度$的优势在于：
即使两个分布没有重叠或者重叠非常少，$Wasserstein距离$仍然能反应两个分布的远近。

4.3 代码实现

# @File    : WGAN_2017.py
import os
import numpy as np
import matplotlib.pyplot as plt
from tensorflow.keras.models import Sequential, Model
from tensorflow.keras.layers import Dense, LeakyReLU, BatchNormalization, Reshape, Input
from tensorflow.keras.layers import UpSampling2D, Conv2D, Activation, ZeroPadding2D, GlobalAveragePooling2D
from tensorflow.keras.layers import Flatten, Dropout
from tensorflow.keras.datasets import mnist
from tensorflow.keras.optimizers import RMSprop
import tensorflow.keras.backend as K


class WGAN():
    def __init__(self):
        self.img_rows = 28
        self.img_cols = 28
        self.channels = 1
        self.img_shape = (self.img_rows, self.img_cols, self.channels)
        self.latent_dim = 100

        self.n_critic = 5
        self.clip_value = 0.01
        optimizer = RMSprop(lr=0.00005)

        self.critic = self.build_critic()
        self.critic.compile(loss=self.wasserstein_loss,
                            optimizer=optimizer,
                            metrics=['accuracy'])

        self.generator = self.build_generator()

        z = Input(shape=(self.latent_dim,))
        img = self.generator(z)

        self.critic.trainable = False

        valid = self.critic(img)

        self.combined = Model(z, valid)
        self.combined.compile(loss=self.wasserstein_loss,
                              optimizer=optimizer,
                              metrics=['accuracy'])

    def build_generator(self):
        model = Sequential()

        model.add(Dense(128 * 7 * 7, activation="relu", input_dim=self.latent_dim))
        model.add(Reshape((7, 7, 128)))
        model.add(UpSampling2D())
        model.add(Conv2D(128, kernel_size=4, padding="same"))
        model.add(BatchNormalization(momentum=0.8))
        model.add(Activation("relu"))
        model.add(UpSampling2D())
        model.add(Conv2D(64, kernel_size=4, padding="same"))
        model.add(BatchNormalization(momentum=0.8))
        model.add(Activation("relu"))
        model.add(Conv2D(self.channels, kernel_size=4, padding="same"))
        model.add(Activation("tanh"))

        model.summary()

        noise = Input(shape=(self.latent_dim,))
        img = model(noise)

        return Model(noise, img)

    def build_critic(self):
        model = Sequential()

        model.add(Conv2D(16, kernel_size=3, strides=2, input_shape=self.img_shape, padding="same"))
        model.add(LeakyReLU(alpha=0.2))
        model.add(Dropout(0.25))
        model.add(Conv2D(32, kernel_size=3, strides=2, padding="same"))
        model.add(ZeroPadding2D(padding=((0,1),(0,1))))
        model.add(BatchNormalization(momentum=0.8))
        model.add(LeakyReLU(alpha=0.2))
        model.add(Dropout(0.25))
        model.add(Conv2D(64, kernel_size=3, strides=2, padding="same"))
        model.add(BatchNormalization(momentum=0.8))
        model.add(LeakyReLU(alpha=0.2))
        model.add(Dropout(0.25))
        model.add(Conv2D(128, kernel_size=3, strides=1, padding="same"))
        model.add(BatchNormalization(momentum=0.8))
        model.add(LeakyReLU(alpha=0.2))
        model.add(Dropout(0.25))
        model.add(Flatten())
        model.add(Dense(1))

        model.summary()

        img = Input(shape=self.img_shape)
        validity = model(img)

        return Model(img, validity)

    def wasserstein_loss(self, y_true, y_pred):
        return K.mean(y_true * y_pred)

    def train(self, epochs, batch_size=128, sample_interval=50):

        (X_train, _), (_, _) = mnist.load_data()

        X_train = (X_train.astype(np.float32) - 127.5) / 127.5
        X_train = np.expand_dims(X_train, axis=3)

        valid = -np.ones((batch_size, 1))
        fake = np.ones((batch_size, 1))

        for epoch in range(epochs):

            for _ in range(self.n_critic):

                idx = np.random.randint(0, X_train.shape[0], batch_size)
                imgs = X_train[idx]

                noise = np.random.normal(0, 1, (batch_size, self.latent_dim))

                gen_imgs = self.generator.predict(noise)

                d_loss_real = self.critic.train_on_batch(imgs, valid)
                d_loss_fake = self.critic.train_on_batch(gen_imgs, fake)
                d_loss = 0.5 * np.add(d_loss_fake, d_loss_real)

                for l in self.critic.layers:
                    weights = l.get_weights()
                    weights = [np.clip(w, -self.clip_value, self.clip_value) for w in weights]
                    l.set_weights(weights)

            g_loss = self.combined.train_on_batch(noise, valid)

            print("%d [D loss: %f] [G loss: %f]" % (epoch, 1 - d_loss[0], 1 - g_loss[0]))

            if epoch % sample_interval == 0:
                self.sample_images(epoch)

    def sample_images(self, epoch):
        r, c = 5, 5
        noise = np.random.normal(0, 1, (r * c, self.latent_dim))
        gen_imgs = self.generator.predict(noise)

        gen_imgs = 0.5 * gen_imgs + 0.5

        fig, axs = plt.subplots(r, c)
        cnt = 0
        for i in range(r):
            for j in range(c):
                axs[i,j].imshow(gen_imgs[cnt, :,:,0], cmap='gray')
                axs[i,j].axis('off')
                cnt += 1
        fig.savefig("images/mnist_%d.png" % epoch)
        plt.close()


if __name__ == '__main__':
    wgan = WGAN()
    wgan.train(epochs=20000, batch_size=32, sample_interval=100)

tree

test
│  WGAN_2017.py
└─ images