深度生成模型

import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from torchvision import transforms
from torchvision.utils import save_image
import matplotlib.pyplot as plt
import numpy as np

# 配置中文字体支持
plt.rcParams['font.sans-serif'] = [
    'Noto Sans CJK SC', 'Noto Sans CJK JP', 'SimHei',
    'Microsoft YaHei', 'WenQuanYi Micro Hei', 'DejaVu Sans'
]
plt.rcParams['axes.unicode_minus'] = False

# 设备配置
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 如果不存在则创建目录
sample_dir = '../../../data/vae_samples'
if not os.path.exists(sample_dir):
    os.makedirs(sample_dir)

# 超参数
image_size = 784
h_dim = 400
z_dim = 20
num_epochs = 15
batch_size = 128
learning_rate = 1e-3

# MNIST 数据集
dataset = torchvision.datasets.MNIST(root='../../../dataset/mnist',
                                     train=True,
                                     transform=transforms.ToTensor(),
                                     download=True)

# 数据加载器
data_loader = torch.utils.data.DataLoader(dataset=dataset,
                                          batch_size=batch_size,
                                          shuffle=True)

# VAE 模型
class VAE(nn.Module):
    def __init__(self, image_size=784, h_dim=400, z_dim=20):
        super(VAE, self).__init__()
        self.fc1 = nn.Linear(image_size, h_dim)
        self.fc2 = nn.Linear(h_dim, z_dim)
        self.fc3 = nn.Linear(h_dim, z_dim)
        self.fc4 = nn.Linear(z_dim, h_dim)
        self.fc5 = nn.Linear(h_dim, image_size)

    def encode(self, x):
        h = F.relu(self.fc1(x))
        return self.fc2(h), self.fc3(h)

    def reparameterize(self, mu, log_var):
        std = torch.exp(0.5*log_var)
        eps = torch.randn_like(std)
        return mu + eps*std

    def decode(self, z):
        h = F.relu(self.fc4(z))
        return torch.sigmoid(self.fc5(h))

    def forward(self, x):
        mu, log_var = self.encode(x.view(-1, image_size))
        z = self.reparameterize(mu, log_var)
        x_reconst = self.decode(z)
        return x_reconst, mu, log_var

model = VAE().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

# 损失函数
def loss_function(recon_x, x, mu, log_var):
    BCE = F.binary_cross_entropy(recon_x, x.view(-1, 784), reduction='sum')
    KLD = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())
    return BCE + KLD

# 开始训练
for epoch in range(num_epochs):
    for i, (x, _) in enumerate(data_loader):
        x = x.to(device)
        x_reconst, mu, log_var = model(x)

        loss = loss_function(x_reconst, x, mu, log_var)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if (i+1) % 100 == 0:
            print ("Epoch[{}/{}], Step [{}/{}], Reconst Loss: {:.4f}"
                   .format(epoch+1, num_epochs, i+1, len(data_loader), loss.item()/len(x)))

# 保存并显示图像
with torch.no_grad():
    # 保存采样图像
    z = torch.randn(batch_size, z_dim).to(device)
    out = model.decode(z).view(-1, 1, 28, 28)
    save_image(out, os.path.join(sample_dir, f'sampled-{epoch+1}.png'))

    # 保存重构图像
    out, _, _ = model(x)
    x_concat = torch.cat([x.view(-1, 1, 28, 28), out.view(-1, 1, 28, 28)], dim=3)
    save_image(x_concat, os.path.join(sample_dir, f'reconst-{epoch+1}.png'))
    
    # 可视化重构图像
    fig = plt.figure(figsize=(10, 2))
    plt.suptitle("重构图像 (奇数: 原始, 偶数: 重构)")
    for i in range(10):
        ax = fig.add_subplot(2, 10, i + 1)
        ax.imshow(x[i].cpu().numpy().reshape(28, 28), cmap='gray')
        ax.axis('off')
        ax = fig.add_subplot(2, 10, i + 11)
        ax.imshow(out[i].cpu().numpy().reshape(28, 28), cmap='gray')
        ax.axis('off')
    plt.show()

    # 可视化生成图像
    fig = plt.figure(figsize=(10, 2))
    plt.suptitle("生成图像")
    z = torch.randn(20, z_dim).to(device)
    generated = model.decode(z)
    for i in range(20):
        ax = fig.add_subplot(2, 10, i + 1)
        ax.imshow(generated[i].cpu().numpy().reshape(28, 28), cmap='gray')
        ax.axis('off')
    plt.show()

import os
import torch
import torch.nn as nn
import torchvision
from torchvision import transforms
from torchvision.utils import save_image
import matplotlib.pyplot as plt
import numpy as np

# 配置中文字体支持
plt.rcParams['font.sans-serif'] = [
    'Noto Sans CJK SC', 'Noto Sans CJK JP', 'SimHei',
    'Microsoft YaHei', 'WenQuanYi Micro Hei', 'DejaVu Sans'
]
plt.rcParams['axes.unicode_minus'] = False

# 设备配置
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 如果不存在则创建目录
sample_dir = '../../../data/gan_samples'
if not os.path.exists(sample_dir):
    os.makedirs(sample_dir)

# 超参数
latent_size = 64
hidden_size = 256
image_size = 784
num_epochs = 200
batch_size = 100
learning_rate = 0.0002

# 图像处理
transform = transforms.Compose([
                transforms.ToTensor(),
                transforms.Normalize(mean=(0.5,),
                                     std=(0.5,))])

# MNIST 数据集
mnist = torchvision.datasets.MNIST(root='../../../dataset/mnist',
                                   train=True,
                                   transform=transform,
                                   download=True)

# 数据加载器
data_loader = torch.utils.data.DataLoader(dataset=mnist,
                                          batch_size=batch_size, 
                                          shuffle=True)

# 判别器
D = nn.Sequential(
    nn.Linear(image_size, hidden_size),
    nn.LeakyReLU(0.2),
    nn.Linear(hidden_size, hidden_size),
    nn.LeakyReLU(0.2),
    nn.Linear(hidden_size, 1),
    nn.Sigmoid())

# 生成器 
G = nn.Sequential(
    nn.Linear(latent_size, hidden_size),
    nn.ReLU(),
    nn.Linear(hidden_size, hidden_size),
    nn.ReLU(),
    nn.Linear(hidden_size, image_size),
    nn.Tanh())

# 设备设置
D = D.to(device)
G = G.to(device)

# 二元交叉熵损失和优化器
criterion = nn.BCELoss()
d_optimizer = torch.optim.Adam(D.parameters(), lr=learning_rate)
g_optimizer = torch.optim.Adam(G.parameters(), lr=learning_rate)

def denorm(x):
    out = (x + 1) / 2
    return out.clamp(0, 1)

def reset_grad():
    d_optimizer.zero_grad()
    g_optimizer.zero_grad()

# 开始训练
total_step = len(data_loader)
for epoch in range(num_epochs):
    for i, (images, _) in enumerate(data_loader):
        images = images.reshape(batch_size, -1).to(device)
        
        # 定义图像是真实还是虚假的标签
        real_labels = torch.ones(batch_size, 1).to(device)
        fake_labels = torch.zeros(batch_size, 1).to(device)

        # ================================================================== #
        #                      训练判别器                                    #
        # ================================================================== #

        # 使用真实图像计算 BCE 损失
        outputs = D(images)
        d_loss_real = criterion(outputs, real_labels)
        real_score = outputs
        
        # 使用生成图像（由潜在空间点生成）计算 BCE 损失
        z = torch.randn(batch_size, latent_size).to(device)
        fake_images = G(z)
        outputs = D(fake_images)
        d_loss_fake = criterion(outputs, fake_labels)
        fake_score = outputs
        
        # 判别器总损失
        d_loss = d_loss_real + d_loss_fake
        
        # 清除梯度并反向传播
        reset_grad()
        d_loss.backward()
        d_optimizer.step()
        
        # ================================================================== #
        #                        训练生成器                                  #
        # ================================================================== #

        # 非饱和生成器损失：我们希望生成器欺骗判别器，
        # 因此使用 real_labels 作为目标。当 D(G(z)) 接近 0 时，
        # 这比极小极大目标提供更好的梯度。
        z = torch.randn(batch_size, latent_size).to(device)
        fake_images = G(z)
        outputs = D(fake_images)
        g_loss = criterion(outputs, real_labels)
        
        # 清除梯度并反向传播
        reset_grad()
        g_loss.backward()
        g_optimizer.step()
        
        if (i+1) % 200 == 0:
            print('Epoch [{}/{}], Step [{}/{}], d_loss: {:.4f}, g_loss: {:.4f}, D(x): {:.2f}, D(G(z)): {:.2f}'
                  .format(epoch, num_epochs, i+1, total_step, d_loss.item(), g_loss.item(), 
                          real_score.mean().item(), fake_score.mean().item()))
    
    # 保存真实图像
    if (epoch+1) == 1:
        images = images.reshape(images.size(0), 1, 28, 28)
        save_image(denorm(images), os.path.join(sample_dir, 'real_images.png'))
    
    # 保存采样图像
    fake_images = fake_images.reshape(fake_images.size(0), 1, 28, 28)
    save_image(denorm(fake_images), os.path.join(sample_dir, 'fake_images-{}.png'.format(epoch+1)))

# 保存模型检查点
torch.save(G.state_dict(), os.path.join(sample_dir, 'G.ckpt'))
torch.save(D.state_dict(), os.path.join(sample_dir, 'D.ckpt'))

# 可视化
real_images_path = os.path.join(sample_dir, 'real_images.png')
fake_images_path = os.path.join(sample_dir, f'fake_images-{num_epochs}.png')

fig = plt.figure(figsize=(10, 5))
plt.subplot(1,2,1)
plt.title("真实图像");
plt.axis('off')
plt.imshow(np.transpose(torchvision.io.read_image(real_images_path),(1,2,0)))

plt.subplot(1,2,2)
plt.title("伪造图像")
plt.axis('off')
plt.imshow(np.transpose(torchvision.io.read_image(fake_images_path),(1,2,0)))

plt.show()