本文主要是介绍AE——重构数字(Pytorch+mnist),希望对大家解决编程问题提供一定的参考价值,需要的开发者们随着小编来一起学习吧!
1、简介
- AE(自编码器)由编码器和解码器组成,编码器将输入数据映射到潜在空间,解码器将潜在表示映射回原始输入空间。
- AE的训练目标通常是最小化重构误差,即尽可能地重构输入数据,使得解码器输出与原始输入尽可能接近。
- AE通常用于数据压缩、去噪、特征提取等任务。
- 本文利用AE,输入数字图像。训练后,输入测试数字图像,重构生成新的数字图像。
- 【注】本文案例需要输入才能生成输出,目标是重构,而不是生成。
- 可以看出,重构图片和原始图片差别不大。
2、代码
-
import matplotlib.pyplot as plt import torch import torch.nn as nn import torch.optim as optim import torchvision# 设置种子和其他配置 seed = 42 # 设置随机种子 torch.manual_seed(seed) torch.backends.cudnn.benchmark = False # 禁用 cuDNN 的自动寻找最佳算法 torch.backends.cudnn.deterministic = True # 设置 cuDNN 为确定性模式# 设置批大小、学习周期和学习率 batch_size = 512 epochs = 30 learning_rate = 1e-3# 载入 MNIST 数据集中的图片进行训练 transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor()]) # 将图像转换为张量train_dataset = torchvision.datasets.MNIST(root="~/torch_datasets", train=True, transform=transform, download=True ) # 加载 MNIST 数据集的训练集,设置路径、转换和下载为 Truetrain_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True ) # 创建一个数据加载器,用于加载训练数据,设置批处理大小和是否随机打乱数据# 在一个类中编写编码器和解码器层。为编码器和解码器层的组件都定义了全连接层 class AE(nn.Module):def __init__(self, **kwargs):super().__init__()self.encoder_hidden_layer = nn.Linear(in_features=kwargs["input_shape"], out_features=128) # 编码器隐藏层self.encoder_output_layer = nn.Linear(in_features=128, out_features=128) # 编码器输出层self.decoder_hidden_layer = nn.Linear(in_features=128, out_features=128) # 解码器隐藏层self.decoder_output_layer = nn.Linear(in_features=128, out_features=kwargs["input_shape"]) # 解码器输出层# 定义了模型的前向传播过程,包括激活函数的应用和重构图像的生成def forward(self, features):activation = self.encoder_hidden_layer(features)activation = torch.relu(activation) # ReLU 激活函数,得到编码器的激活值code = self.encoder_output_layer(activation)code = torch.sigmoid(code) # Sigmoid 激活函数,以确保编码后的表示在 [0, 1] 范围内activation = self.decoder_hidden_layer(code)activation = torch.relu(activation)activation = self.decoder_output_layer(activation)reconstructed = torch.sigmoid(activation)return reconstructed# 在使用定义的 AE 类之前,有以下事情要做: # 配置要在哪个设备上运行 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")# 建立 AE 模型并载入到 CPU 设备 model = AE(input_shape=784).to(device)# Adam 优化器,学习率 10e-3 optimizer = optim.Adam(model.parameters(), lr=learning_rate)# 使用均方误差(MSE)损失函数 criterion = nn.MSELoss()# 在CPU设备上运行,实例化一个输入大小为784的AE自编码器,并用Adam作为训练优化器用MSELoss作为损失函数 # 训练: for epoch in range(epochs):loss = 0for batch_features, _ in train_loader:# 将小批数据变形为 [N, 784] 矩阵,并加载到 CPU 设备batch_features = batch_features.view(-1, 784).to(device)# 梯度设置为 0,因为 torch 会累加梯度optimizer.zero_grad()# 计算重构outputs = model(batch_features)# 计算训练重建损失train_loss = criterion(outputs, batch_features)# 计算累积梯度train_loss.backward()# 根据当前梯度更新参数optimizer.step()# 将小批量训练损失加到周期损失中loss += train_loss.item()# 计算每个周期的训练损失loss = loss / len(train_loader)# 显示每个周期的训练损失print("epoch : {}/{}, recon loss = {:.8f}".format(epoch + 1, epochs, loss))if __name__ == '__main__':# 用训练过的自编码器提取一些测试用例来重构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=10, 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 = 10 # 设置要显示的图像数量plt.figure(figsize=(20, 4)) # 创建一个新的 Matplotlib 图形,设置图形大小为 (20, 4)for index in range(number): # 遍历要显示的图像数量# 显示原始图ax = plt.subplot(2, number, index + 1)plt.imshow(test_examples[index].cpu().numpy().reshape(28, 28))plt.gray()ax.get_xaxis().set_visible(False)ax.get_yaxis().set_visible(False)# 显示重构图ax = plt.subplot(2, number, index + 1 + number)plt.imshow(reconstruction[index].cpu().numpy().reshape(28, 28))plt.gray()ax.get_xaxis().set_visible(False)ax.get_yaxis().set_visible(False)plt.savefig('reconstruction_results.png') # 保存图像plt.show()
这篇关于AE——重构数字(Pytorch+mnist)的文章就介绍到这儿,希望我们推荐的文章对编程师们有所帮助!