import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torch.nn.functional as F
from torchvision.utils import save_image
import matplotlib.pyplot as plt
# 变分自编码器
class VAE(nn.Module):
def __init__(self):
super(VAE, self).__init__()
# 编码器层
self.fc1 = nn.Linear(input_size, 512) # 编码器输入层
self.fc2 = nn.Linear(512, latent_size)
self.fc3 = nn.Linear(512, latent_size)
# 解码器层
self.fc4 = nn.Linear(latent_size, 512) # 解码器输入层
self.fc5 = nn.Linear(512, input_size) # 解码器输出层
# 编码器部分
def encode(self, x):
x = F.relu(self.fc1(x)) # 编码器的隐藏表示
mu = self.fc2(x) # 潜在空间均值
log_var = self.fc3(x) # 潜在空间对数方差
return mu, log_var
# 重参数化技巧
def reparameterize(self, mu, log_var): # 从编码器输出的均值和对数方差中采样得到潜在变量z
std = torch.exp(0.5 * log_var) # 计算标准差
eps = torch.randn_like(std) # 从标准正态分布中采样得到随机噪声
return mu + eps * std # 根据重参数化公式计算潜在变量z
# 解码器部分
def decode(self, z):
z = F.relu(self.fc4(z)) # 将潜在变量 z 解码为重构图像
return torch.sigmoid(self.fc5(z)) # 将隐藏表示映射回输入图像大小,并应用 sigmoid 激活函数,以产生重构图像
# 前向传播
def forward(self, x): # 输入图像 x 通过编码器和解码器,得到重构图像和潜在变量的均值和对数方差
mu, log_var = self.encode(x.view(-1, input_size))
z = self.reparameterize(mu, log_var)
return self.decode(z), mu, log_var
# 使用重构损失和 KL 散度作为损失函数
def loss_function(recon_x, x, mu, log_var): # 参数:重构的图像、原始图像、潜在变量的均值、潜在变量的对数方差
MSE = F.mse_loss(recon_x, x.view(-1, input_size), reduction='sum') # 计算重构图像 recon_x 和原始图像 x 之间的均方误差
KLD = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp()) # 计算潜在变量的KL散度
return MSE + KLD # 返回二进制交叉熵损失和 KLD 损失的总和作为最终的损失值
if __name__ == '__main__':
batch_size = 512 # 批次大小
epochs = 30 # 学习周期
sample_interval = 10 # 保存结果的周期
learning_rate = 0.001 # 学习率
input_size = 784 # 输入大小
latent_size = 64 # 潜在空间大小
# 载入 MNIST 数据集中的图片进行训练
transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor()]) # 将图像转换为张量
train_dataset = torchvision.datasets.MNIST(
root="~/torch_datasets", train=True, transform=transform, download=True
) # 加载 MNIST 数据集的训练集,设置路径、转换和下载为 True
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=batch_size, shuffle=True
) # 创建一个数据加载器,用于加载训练数据,设置批处理大小和是否随机打乱数据
# 在使用定义的 AE 类之前,有以下事情要做:
# 配置要在哪个设备上运行
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 建立 VAE 模型并载入到 CPU 设备
model = VAE().to(device)
# Adam 优化器,学习率
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
# 训练
for epoch in range(epochs):
train_loss = 0
for batch_idx, (data, _) in enumerate(train_loader):
data = data.to(device) # 将输入数据移动到设备(GPU 或 CPU)上
optimizer.zero_grad() # 进行反向传播之前,需要将优化器中的梯度清零,以避免梯度的累积
# 重构图像 recon_batch、潜在变量的均值 mu 和对数方差 log_var
recon_batch, mu, log_var = model(data)
loss = loss_function(recon_batch, data, mu, log_var) # 计算损失
loss.backward() # 计算损失相对于模型参数的梯度
train_loss += loss.item()
optimizer.step() # 更新模型参数
train_loss = train_loss / len(train_loader) # # 计算每个周期的训练损失
print('Epoch [{}/{}], Loss: {:.3f}'.format(epoch + 1, epochs, train_loss))
# 用训练过的自编码器提取一些测试用例来重构
test_dataset = torchvision.datasets.MNIST(
root="~/torch_datasets", train=False, transform=transform, download=True
) # 加载 MNIST 测试数据集
test_loader = torch.utils.data.DataLoader(
test_dataset, batch_size=16, shuffle=False
) # 创建一个测试数据加载器
test_examples = None
# 通过循环遍历测试数据加载器,获取一个批次的图像数据
with torch.no_grad(): # 使用 torch.no_grad() 上下文管理器,确保在该上下文中不会进行梯度计算
for batch_features in test_loader: # 历测试数据加载器中的每个批次的图像数据
batch_features = batch_features[0] # 获取当前批次的图像数据
test_examples = batch_features.view(-1, 784).to(
device) # 将当前批次的图像数据转换为大小为 (批大小, 784) 的张量,并加载到指定的设备(CPU 或 GPU)上
reconstruction, _, _ = model(test_examples) # 使用训练好的自编码器模型对测试数据进行重构,即生成重构的图像
break
# 试着用训练过的自编码器重建一些测试图像
with torch.no_grad():
number = 16 # 设置要显示的图像数量
plt.figure(figsize=(60, 10)) # 创建一个新的 Matplotlib 图形,设置图形大小为 (16, 16)
for index in range(number): # 遍历要显示的图像数量
# 显示原始图
ax = plt.subplot(2, number, index + 1) # 行 列 位置
plt.imshow(test_examples[index].cpu().numpy().reshape(28, 28))
plt.gray()
plt.axis('off')
# 显示重构图
ax = plt.subplot(2, number, index + 1 + number)
plt.imshow(reconstruction[index].cpu().numpy().reshape(28, 28))
plt.gray()
plt.axis('off')
plt.savefig('reconstruction_results.png') # 保存图像
plt.show()