import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim from torch.utils.data import DataLoader import torchvision.transforms as transforms from torchvision.datasets import MNIST from torchvision.utils import save_image import numpy as np import matplotlib.pyplot as plt from tqdm import tqdm # 设置随机种子和设备 torch.manual_seed(42) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 超参数 batch_size = 128 epochs = 30 learning_rate = 1e-3 latent_dim = 20 # 潜在空间维度 img_size = 28 num_classes = 10 # 数据加载 transform = transforms.Compose([ transforms.ToTensor(), ]) train_dataset = MNIST(root='./data', train=True, transform=transform, download=True) test_dataset = MNIST(root='./data', train=False, transform=transform, download=True) train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False) # 条件变分自编码器模型 class CVAE(nn.Module): def __init__(self, latent_dim, num_classes): super(CVAE, self).__init__() # 图像大小和条件信息 self.img_size = 28 self.latent_dim = latent_dim self.num_classes = num_classes # 编码器网络 self.encoder = nn.Sequential( nn.Conv2d(1 + num_classes, 32, kernel_size=3, stride=2, padding=1), # 14x14 nn.LeakyReLU(0.2), nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1), # 7x7 nn.LeakyReLU(0.2), nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1), # 7x7 nn.LeakyReLU(0.2), nn.Flatten() ) # 计算展平后的特征大小 self.flat_size = 128 * 7 * 7 # 均值和方差预测层 self.fc_mu = nn.Linear(self.flat_size, latent_dim) self.fc_logvar = nn.Linear(self.flat_size, latent_dim) # 解码器输入层 self.decoder_input = nn.Linear(latent_dim + num_classes, 128 * 7 * 7) # 解码器网络 self.decoder = nn.Sequential( nn.Unflatten(1, (128, 7, 7)), nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1), # 14x14 nn.LeakyReLU(0.2), nn.ConvTranspose2d(64, 32, kernel_size=4, stride=2, padding=1), # 28x28 nn.LeakyReLU(0.2), nn.Conv2d(32, 1, kernel_size=3, stride=1, padding=1), nn.Sigmoid() # 输出像素值在[0,1]之间 ) def encode(self, x, c): # 将条件信息嵌入到输入中 c = c.view(-1, self.num_classes, 1, 1).expand(-1, -1, self.img_size, self.img_size) x_c = torch.cat([x, c], dim=1) # 在通道维度上拼接 # 编码器前向传播 h = self.encoder(x_c) mu = self.fc_mu(h) logvar = self.fc_logvar(h) return mu, logvar def reparameterize(self, mu, logvar): # 重参数化技巧 std = torch.exp(0.5 * logvar) eps = torch.randn_like(std) z = mu + eps * std return z def decode(self, z, c): # 将条件信息与潜在表示拼接 z_c = torch.cat([z, c], dim=1) # 解码器前向传播 h = self.decoder_input(z_c) x_recon = self.decoder(h.view(-1, 128, 7, 7)) return x_recon def forward(self, x, c): # 将类标签转换为one-hot向量 c_onehot = F.one_hot(c, self.num_classes).float() # 编码 mu, logvar = self.encode(x, c_onehot) # 采样潜在表示 z = self.reparameterize(mu, logvar) # 解码 x_recon = self.decode(z, c_onehot) return x_recon, mu, logvar def sample(self, num_samples, c): """ 给定条件c，生成样本 c: (num_samples,) 类标签 """ # 将类标签转换为one-hot向量 c_onehot = F.one_hot(c, self.num_classes).float() # 从标准正态分布采样潜在向量 z = torch.randn(num_samples, self.latent_dim).to(device) # 解码生成样本 samples = self.decode(z, c_onehot) return samples # 损失函数 def loss_function(recon_x, x, mu, logvar, beta=1.0): # 重建损失 (二值交叉熵) BCE = F.binary_cross_entropy(recon_x, x, reduction='sum') # KL散度 KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp()) return BCE + beta * KLD # 实例化模型 model = CVAE(latent_dim=latent_dim, num_classes=num_classes).to(device) optimizer = optim.Adam(model.parameters(), lr=learning_rate) # 训练函数 def train(epoch): model.train() train_loss = 0 for batch_idx, (data, target) in enumerate(tqdm(train_loader)): data, target = data.to(device), target.to(device) # 前向传播 optimizer.zero_grad() recon_batch, mu, logvar = model(data, target) # 计算损失 loss = loss_function(recon_batch, data, mu, logvar) # 反向传播 loss.backward() train_loss += loss.item() optimizer.step() print(f'====> Epoch: {epoch} Average loss: {train_loss / len(train_loader.dataset):.4f}') # 测试函数 def test(epoch): model.eval() test_loss = 0 with torch.no_grad(): for data, target in test_loader: data, target = data.to(device), target.to(device) # 前向传播 recon_batch, mu, logvar = model(data, target) # 计算损失 test_loss += loss_function(recon_batch, data, mu, logvar).item() test_loss /= len(test_loader.dataset) print(f'====> Test set loss: {test_loss:.4f}') return test_loss # 生成条件样本并可视化 def generate_samples(epoch): model.eval() with torch.no_grad(): # 为每个数字类别生成样本 n = 10 # 每个类别生成的样本数 digit_size = 28 figure = np.zeros((digit_size * n, digit_size * num_classes)) # 对每个类别生成n个样本 for i in range(num_classes): # 创建类标签 c = torch.tensor([i] * n).to(device) # 生成样本 samples = model.sample(n, c) samples = samples.cpu().numpy() # 将生成的样本放入图中对应位置 for j in range(n): figure[j * digit_size: (j + 1) * digit_size, i * digit_size: (i + 1) * digit_size] = samples[j, 0] # 保存生成的样本图 plt.figure(figsize=(10, 10)) plt.imshow(figure, cmap='gray') plt.axis('off') plt.tight_layout() plt.savefig(f'cvae_samples_epoch_{epoch}.png') plt.close() # 潜在空间可视化 def visualize_latent_space(epoch): model.eval() with torch.no_grad(): # 获取测试数据 test_data = [] test_labels = [] for data, label in test_loader: if len(test_data) < 1000: # 仅使用1000个样本进行可视化 test_data.append(data) test_labels.append(label) else: break test_data = torch.cat(test_data, dim=0)[:1000].to(device) test_labels = torch.cat(test_labels, dim=0)[:1000].to(device) # 获取潜在表示 c_onehot = F.one_hot(test_labels, num_classes).float() mu, _ = model.encode(test_data, c_onehot) z = mu.cpu().numpy() # 使用PCA或t-SNE降维可视化 if latent_dim > 2: from sklearn.decomposition import PCA pca = PCA(n_components=2) z_2d = pca.fit_transform(z) else: z_2d = z # 绘制散点图 plt.figure(figsize=(10, 8)) scatter = plt.scatter(z_2d[:, 0], z_2d[:, 1], c=test_labels.cpu(), cmap='tab10') plt.colorbar(scatter) plt.title(f'Latent Space (Epoch {epoch})') plt.xlabel('Dimension 1') plt.ylabel('Dimension 2') plt.tight_layout() plt.savefig(f'latent_space_epoch_{epoch}.png') plt.close() # 插值实验 def generate_interpolations(): model.eval() with torch.no_grad(): # 选择两个数字进行插值 start_digit = 3 end_digit = 7 num_steps = 10 # 创建起始和结束条件向量 c_start = F.one_hot(torch.tensor([start_digit]), num_classes).float().to(device) c_end = F.one_hot(torch.tensor([end_digit]), num_classes).float().to(device) # 从潜在空间采样一个向量 z = torch.randn(1, latent_dim).to(device) # 创建条件插值 figure = np.zeros((digit_size, digit_size * num_steps)) for i in range(num_steps): # 线性插值条件向量 alpha = i / (num_steps - 1) c_interp = c_start * (1 - alpha) + c_end * alpha # 生成样本 sample = model.decode(z, c_interp) sample = sample.cpu().numpy() # 将生成的样本放入图中 figure[:, i * digit_size: (i + 1) * digit_size] = sample[0, 0] # 保存插值结果 plt.figure(figsize=(12, 2)) plt.imshow(figure, cmap='gray') plt.axis('off') plt.title(f'Interpolation from {start_digit} to {end_digit}') plt.tight_layout() plt.savefig(f'cvae_interpolation_{start_digit}_to_{end_digit}.png') plt.close() # 训练模型 best_loss = float('inf') for epoch in range(1, epochs + 1): train(epoch) test_loss = test(epoch) # 每隔5个epoch生成样本 if epoch % 5 == 0 or epoch == 1: generate_samples(epoch) visualize_latent_space(epoch) # 保存最佳模型 if test_loss < best_loss: best_loss = test_loss torch.save(model.state_dict(), 'cvae_mnist.pth') # 加载最佳模型并进行插值实验 model.load_state_dict(torch.load('cvae_mnist.pth')) generate_interpolations() # 条件生成实验：根据用户输入的数字生成样本 def generate_requested_digits(): model.eval() with torch.no_grad(): # 为每个数字生成一排样本 samples_per_digit = 8 figure = np.zeros((digit_size * num_classes, digit_size * samples_per_digit)) for i in range(num_classes): # 为当前数字创建条件向量 c = torch.tensor([i] * samples_per_digit).to(device) # 生成样本 samples = model.sample(samples_per_digit, c) samples = samples.cpu().numpy() # 将样本放入图中 for j in range(samples_per_digit): figure[i * digit_size: (i + 1) * digit_size, j * digit_size: (j + 1) * digit_size] = samples[j, 0] # 保存生成的样本 plt.figure(figsize=(samples_per_digit, num_classes)) plt.imshow(figure, cmap='gray') plt.axis('off') plt.title('Generated Digits (0-9)') plt.tight_layout() plt.savefig('cvae_all_digits.png') plt.close() # 执行条件生成实验 generate_requested_digits() print("Training and generation completed!")