PyTorch-Tutorial项目中的自编码器实现详解-CSDN博客

本文链接：https://blog.csdn.net/gitblog_00612/article/details/148440905

PyTorch-Tutorial项目中的自编码器实现详解

PyTorch-Tutorial Build your neural network easy and fast, 莫烦Python中文教学项目地址: https://gitcode.com/gh_mirrors/pyt/PyTorch-Tutorial

自编码器概述

自编码器(AutoEncoder)是一种无监督学习的神经网络模型，主要用于数据的降维和特征提取。它通过将输入数据压缩到一个低维空间(编码)，然后再从这个低维表示重建原始数据(解码)，从而学习数据的有用特征。

代码实现解析

1. 数据准备

首先加载MNIST手写数字数据集，这是深度学习中最常用的基准数据集之一：

train_data = torchvision.datasets.MNIST(
    root='./mnist/',
    train=True,
    transform=torchvision.transforms.ToTensor(),
    download=DOWNLOAD_MNIST
)

数据加载器将数据集分成小批量，便于训练：

train_loader = Data.DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True)

2. 自编码器模型结构

自编码器由编码器和解码器两部分组成：

class AutoEncoder(nn.Module):
    def __init__(self):
        super(AutoEncoder, self).__init__()
        
        # 编码器部分
        self.encoder = nn.Sequential(
            nn.Linear(28*28, 128),
            nn.Tanh(),
            nn.Linear(128, 64),
            nn.Tanh(),
            nn.Linear(64, 12),
            nn.Tanh(),
            nn.Linear(12, 3),  # 压缩到3维特征
        )
        
        # 解码器部分
        self.decoder = nn.Sequential(
            nn.Linear(3, 12),
            nn.Tanh(),
            nn.Linear(12, 64),
            nn.Tanh(),
            nn.Linear(64, 128),
            nn.Tanh(),
            nn.Linear(128, 28*28),
            nn.Sigmoid()  # 输出在0-1范围内
        )

编码器将784维(28×28)的输入图像逐步压缩到3维空间，解码器则从3维空间重建原始图像。使用Tanh作为激活函数，最后解码器输出使用Sigmoid确保像素值在0-1之间。

3. 训练过程

训练过程使用均方误差(MSE)作为损失函数，Adam优化器进行参数更新：

optimizer = torch.optim.Adam(autoencoder.parameters(), lr=LR)
loss_func = nn.MSELoss()

for epoch in range(EPOCH):
    for step, (x, b_label) in enumerate(train_loader):
        b_x = x.view(-1, 28*28)  # 展平输入
        b_y = x.view(-1, 28*28)  # 目标输出
        
        encoded, decoded = autoencoder(b_x)
        loss = loss_func(decoded, b_y)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

4. 可视化结果

训练过程中实时显示原始图像和重建图像的对比：

# 初始化图像显示
f, a = plt.subplots(2, N_TEST_IMG, figsize=(5, 2))
plt.ion()

# 训练过程中更新显示
_, decoded_data = autoencoder(view_data)
for i in range(N_TEST_IMG):
    a[1][i].imshow(np.reshape(decoded_data.data.numpy()[i], (28, 28)), cmap='gray')

训练完成后，将200个样本的3维编码结果在3D空间中可视化：

fig = plt.figure(2); ax = Axes3D(fig)
X, Y, Z = encoded_data.data[:, 0].numpy(), encoded_data.data[:, 1].numpy(), encoded_data.data[:, 2].numpy()
values = train_data.train_labels[:200].numpy()
for x, y, z, s in zip(X, Y, Z, values):
    c = cm.rainbow(int(255*s/9)); ax.text(x, y, z, s, backgroundcolor=c)