【PyTorch][chapter 18][李宏毅深度学习]【无监督学习][ VAE]

 # -*- coding: utf-8 -*-

    
 """
    
 Created on Mon Feb 26 15:47:20 2024
    
   5. @author: chengxf2
    
 """
    
  
    
 import torch
    
 import torch.nn as nn
    
 import torch.optim as optim
    
 import torchvision
    
 import torchvision.transforms as transforms   # transforms用于数据预处理
    
  
    
  
    
 # 定义变分自编码器（VAE）模型
    
 class VAE(nn.Module):
    
     def __init__(self, latent_dim):
    
     super(VAE, self).__init__()
    
     
    
     # Encoder
    
     self.encoder = nn.Sequential(
    
         nn.Linear(in_features=784, out_features=256),
    
         nn.ReLU(),
    
         nn.Linear(in_features=256, out_features=128),
    
         nn.ReLU(),
    
         nn.Linear(in_features=128, out_features=latent_dim*2),  # 输出均值和方差
    
         nn.ReLU()
    
     )
    
     
    
     # Decoder
    
     self.decoder = nn.Sequential(
    
         nn.Linear(in_features =latent_dim , out_features=128),
    
         nn.ReLU(),
    
         nn.Linear(in_features=128, out_features=256),
    
         nn.ReLU(),
    
         nn.Linear(in_features=256, out_features=784),
    
         nn.Sigmoid()
    
     )
    
     
    
     def reparameterize(self, mu, logvar):
    
     
    
     std = torch.exp(logvar/2.0)  # 计算标准差,Encoder 出来的可能有负的值,标准差为非负值,所以要乘以exp
    
     eps = torch.randn_like(std)  # 从标准正态分布中采样噪声
    
     z = mu + eps * std  # 重参数化技巧
    
     return z
    
     
    
     def forward(self, x):
    
     # 编码[batch, latent_dim*2]
    
     encoded = self.encoder(x)
    
     #[ z = mu|logvar]
    
     mu, logvar = torch.chunk(encoded, 2, dim=1)  # 将输出分割为均值和方差
    
       
    
     
    
     z = self.reparameterize(mu, logvar)  # 重参数化
    
     
    
     # 解码
    
     decoded = self.decoder(z)
    
     return decoded, mu, logvar
    
  
    
 # 定义训练函数
    
 def train_vae(model, train_loader, num_epochs, learning_rate):
    
     criterion = nn.BCELoss()  # 二元交叉熵损失函数
    
     optimizer = optim.Adam(model.parameters(), lr=learning_rate)  # Adam优化器
    
     
    
     model.train()  # 设置模型为训练模式
    
     
    
     for epoch in range(num_epochs):
    
     total_loss = 0.0
    
     
    
     for data in train_loader:
    
         images, _ = data
    
         images = images.view(images.size(0), -1)  # 展平输入图像
    
         
    
         optimizer.zero_grad()
    
         
    
         # 前向传播
    
         outputs, mu, logvar = model(images)
    
         
    
         # 计算重构损失和KL散度
    
         reconstruction_loss = criterion(outputs, images)
    
         kl_divergence = 0.5 * torch.sum( -logvar +mu.pow(2) +logvar.exp()-1)
    
         
    
         # 计算总损失
    
         loss = reconstruction_loss + kl_divergence
    
         
    
         # 反向传播和优化
    
         loss.backward()
    
         optimizer.step()
    
         
    
         total_loss += loss.item()
    
     
    
     # 输出当前训练轮次的损失
    
     print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, num_epochs, total_loss / len(train_loader)))
    
     
    
     print('Training finished.')
    
  
    
 # 示例用法
    
 if __name__ == '__main__':
    
     # 设置超参数
    
   
    
     latent_dim = 32  # 潜在空间维度
    
     num_epochs = 1  # 训练轮次
    
     learning_rate = 1e-4  # 学习率
    
     
    
     # 加载MNIST数据集
    
     train_dataset = torchvision.datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor(), download=True)
    
     train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=100, shuffle=True)
    
     
    
     # 创建VAE模型
    
     model = VAE(latent_dim)
    
     
    
     # 训练VAE模型
    
     train_vae(model, train_loader, num_epochs, learning_rate)