一、WGAN原理

1.1 原始GAN的缺陷

原始GAN通过最小化JS散度（Jensen-Shannon Divergence）训练生成器（Generator）和判别器（Discriminator），但存在两个关键问题：

1. 梯度消失：当真实分布 $P_r$ 和生成分布 $P_g$ 不重叠时，JS散度为常数 $\log 2$，导致梯度为0，无法更新生成器
2. 训练不稳定：判别器容易过拟合，生成器难以收敛

---

1.2 Wasserstein距离的引入

Wasserstein距离（Earth-Mover距离）衡量两个分布 ${\mathbb{P}}_r$（真实分布）和 ${\mathbb{P}}_g$（生成分布）之间的差异：

$W({\mathbb{P}}_r,{\mathbb{P}}_g)={\inf }_{\gamma \in \Pi ({\mathbb{P}}_r,{\mathbb{P}}_g)}{\mathbb{E}}_{(x,y)\sim \gamma }[\Vert x-y\Vert ]$，其中 $\Pi ({\mathbb{P}}_r,{\mathbb{P}}_g)$ 是联合分布集合，边缘分布分别为 ${\mathbb{P}}_r$ 和 ${\mathbb{P}}_g$

关键改进：即使两个分布支撑集不重叠，Wasserstein距离仍能提供有意义的梯度

直观解释：衡量将“概率质量”从 $P_r$ 搬运到 $P_g$ 的最小成本

---

1.3 Kantorovich-Rubinstein对偶

通过对偶形式将问题转化为：

$W({\mathbb{P}}_r,{\mathbb{P}}_g)={\sup }_{\Vert f{\Vert }_L\le 1}{\mathbb{E}}_{x\sim {\mathbb{P}}_r}[f(x)]-{\mathbb{E}}_{x\sim {\mathbb{P}}_g}[f(x)]$，其中 $f$ 是1-Lipschitz函数（满足 $|f(x)-f(y)|\le \Vert x-y\Vert$）

---

1.4 WGAN的优化目标

• 判别器（Critic）：拟合一个1-Lipschitz函数 $f_w$，最大化：
  $L_{\text{critic}}={\mathbb{E}}_{x\sim {\mathbb{P}}_r}[f_w(x)]-{\mathbb{E}}_{z\sim p(z)}[f_w(G_{\theta }(z))]$
• 生成器：最小化Wasserstein距离，即：
  $L_{\text{generator}}=-{\mathbb{E}}_{z\sim p(z)}[f_w(G_{\theta }(z))]$

关键点：

1. 判别器（称为Critic）输出为标量，无需Sigmoid激活
2. 通过权重裁剪（强制参数 $w$) 在 $[-c,c]$ 内）或梯度惩罚（WGAN-GP）近似Lipschitz约束

---

1.4 数学推导步骤

（1）原始Wasserstein距离

$W({\mathbb{P}}_r,{\mathbb{P}}_g)={\inf }_{\gamma }{\mathbb{E}}_{(x,y)\sim \gamma }[\Vert x-y\Vert ]$

（2）对偶形式推导

利用线性规划对偶性，转化为：$W({\mathbb{P}}_r,{\mathbb{P}}_g)={\sup }_{f\in \text{1-Lip}}{\mathbb{E}}_{x\sim {\mathbb{P}}_r}[f(x)]-{\mathbb{E}}_{x\sim {\mathbb{P}}_g}[f(x)]$

（3） 参数化近似
用神经网络 $f_w$ 近似 $f$，优化：${\max }_w{\mathbb{E}}_{x\sim {\mathbb{P}}_r}[f_w(x)]-{\mathbb{E}}_{z\sim p(z)}[f_w(G_{\theta }(z))]$

（4）生成器优化
固定$f_w$，生成器最小化：${\min }_{\theta }-{\mathbb{E}}_{z\sim p(z)}[f_w(G_{\theta }(z))]$

---

1.5 权重裁剪 vs 梯度惩罚

• 权重裁剪（原始WGAN）：
  强制参数 $w$ 在 $[-c,c]$ 内，但可能导致梯度消失或爆炸
• 梯度惩罚（WGAN-GP）：
  添加正则项： $\lambda {\mathbb{E}}_{\hat{x}\sim {\mathbb{P}}_{\hat{x}}}[(\Vert {\nabla }_{\hat{x}}{f}_w(\hat{x}){\Vert }_2-1{)}^2]$，其中 $\hat{x}$ 是真实样本和生成样本的随机插值

---

1.6 优势

1. 训练信号：Critic的损失值与生成样本质量相关（越低表示越真实）
2. 训练稳定性：避免模式崩溃（Mode Collapse）
3. 梯度有意义：即使分布不重叠，仍能提供有效梯度
4. 生成质量高：Wasserstein距离直接反映生成数据与真实数据的差异

1.7 总结

WGAN通过Wasserstein距离的优良性质，解决了传统GAN的训练难题。其数学核心在于对偶形式的转化和Lipschitz约束的实现，后续改进（如WGAN-GP）进一步提升了性能。

---

二、WGAN实现

2.1 导包

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from torchvision import datasets, transforms
from torchvision.utils import save_image
import numpy as npimport os
import time
import matplotlib.pyplot as plt
from tqdm.notebook import tqdm  
from torchsummary import summary# 判断是否存在可用的GPU
device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")# 指定存放日志路径
writer=SummaryWriter(log_dir="./runs/wgan")os.makedirs("./img/wgan_mnist", exist_ok=True) # 存放生成样本目录
os.makedirs("./model", exist_ok=True) # 模型存放目录

2.2 数据加载和处理

# 加载 MNIST 数据集
def load_data(batch_size=64,img_shape=(1,28,28)):transform = transforms.Compose([transforms.ToTensor(),  # 将图像转换为张量transforms.Normalize(mean=[0.5], std=[0.5])  # 归一化到[-1,1]])# 下载训练集和测试集train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)# 创建 DataLoadertrain_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, num_workers=2,shuffle=True)test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, num_workers=2,shuffle=False)return train_loader, test_loader

2.3 构建生成器

class Generator(nn.Module):"""生成器"""def __init__(self, latent_dim=100,img_shape=(1,28,28)):super(Generator,self).__init__()# 网络块def block(in_feat, out_feat, normalize=True):layers = [nn.Linear(in_feat, out_feat)]if normalize:layers.append(nn.BatchNorm1d(out_feat))layers.append(nn.LeakyReLU(negative_slope=0.2, inplace=True))return layersself.model = nn.Sequential(*block(latent_dim, 128, normalize=False),*block(128, 256),*block(256, 512),*block(512, 1024),nn.Linear(1024, int(np.prod(img_shape))),nn.Tanh() # 输出归一化到[-1,1] )def forward(self,z): # 噪声z,2维[batch_size,latent_dim]gen_img=self.model(z) gen_img=gen_img.view(gen_img.shape[0],*img_shape)return gen_img # 4维[batch_size,1,H,W]

2.4 构建判别器

class Discriminator(nn.Module):"""判别器"""def __init__(self,img_shape=(1,28,28)):super(Discriminator, self).__init__()self.model = nn.Sequential(nn.Linear(int(np.prod(img_shape)), 512),nn.LeakyReLU(negative_slope=0.2, inplace=True),nn.Linear(512, 256),nn.LeakyReLU(negative_slope=0.2, inplace=True),nn.Linear(256, 1))def forward(self,img): # 输入图片，4维[batc_size,1,H,W]img=img.view(img.shape[0], -1) pred = self.model(img)return pred # 2维[batch_size,1] 

2.5 训练和保存模型

• WGAN算法流程
  
• 训练和保存

# 设置超参数
batch_size = 64
epochs = 200
lr= 0.00005
latent_dim=100 # 生成器输入噪声向量的长度(维数)
sample_interval=400 #每400次迭代保存生成样本# WGAN的特别设置
num_iter_critic = 5
weight_clip_value = 0.01# 设置图片形状1*28*28
img_shape = (1,28,28)# 加载数据
train_loader,_= load_data(batch_size=batch_size,img_shape=img_shape)# 实例化生成器G、判别器D
G=Generator().to(device)
D=Discriminator().to(device)# 设置优化器
optimizer_G = torch.optim.RMSprop(G.parameters(), lr=lr)
optimizer_D = torch.optim.RMSprop(D.parameters(), lr=lr)# 开始训练
batches_done=0
loader_len=len(train_loader) #训练集加载器的长度
for epoch in range(epochs):# 进入训练模式G.train()D.train()loop = tqdm(train_loader, desc=f"第{epoch+1}轮")for i, (real_imgs, _) in enumerate(loop):real_imgs=real_imgs.to(device)  # [B,C,H,W]# -----------------#  训练判别器# -----------------# 获取噪声样本[B,latent_dim）z=torch.normal(0,1,size=(real_imgs.shape[0],latent_dim),device=device)  #从正态分布中抽样# Step-1 计算判断器损失=判断真实图片损失+判断生成图片损失fake_imgs=G(z).detach()dis_loss=-torch.mean(D(real_imgs)) + torch.mean(D(fake_imgs))# Step-2 更新判别器参数optimizer_D.zero_grad() # 梯度清零dis_loss.backward() #反向传播,计算梯度optimizer_D.step()  #更新判别器 # Step-3 对判别器进行权重裁剪for p in D.parameters():p.data.clamp_(-weight_clip_value,weight_clip_value)# -----------------#  训练生成器# -----------------# 判别器每迭代 num_iter_critic 次,生成器迭代一次if i % num_iter_critic ==0 :gen_imgs=G(z).detach()# 更新生成器参数optimizer_G.zero_grad() #梯度清零gen_loss=-torch.mean(D(gen_imgs))gen_loss.backward() #反向传播,计算梯度optimizer_G.step()  #更新生成器  # 更新进度条loop.set_postfix(gen_loss=f"{gen_loss:.8f}",dis_loss=f"{dis_loss:.8f}")# 每 sample_interval 次迭代保存生成样本if batches_done % sample_interval == 0:save_image(gen_imgs.data[:25], f"./img/wgan_mnist/{epoch}_{i}.png", nrow=5, normalize=True)batches_done += 1print('总共训练用时: %.2f min' % ((time.time() - start_time)/60))#仅保存模型的参数（权重和偏置）,灵活性高，可以在不同的模型结构之间加载参数
torch.save(G.state_dict(), "./model/WGAN_G.pth") 
torch.save(D.state_dict(), "./model/WGAN_D.pth") 

2.6 图片转GIF

from PIL import Imagedef create_gif(img_dir="./img/wgan_mnist", output_file="./img/wgan_mnist/wgan_figure.gif", duration=100):images = []img_paths = [f for f in os.listdir(img_dir) if f.endswith(".png")]# 自定义排序：按 "x_y.png" 的 x 和 y 排序img_paths_sorted = sorted(img_paths,key=lambda x: (int(x.split('_')[0]),  # 第一个数字（如 0_400.png 的 0）int(x.split('_')[1].split('.')[0])  # 第二个数字（如 0_400.png 的 400）))for img_file in img_paths_sorted:img = Image.open(os.path.join(img_dir, img_file))images.append(img)images[0].save(output_file, save_all=True, append_images=images[1:], duration=duration, loop=0)print(f"GIF已保存至 {output_file}")
create_gif()