使用AutoEncoder对MNIST手写数据集进行表征学习（pytorch）

AutoEncoder是一种无监督学习方式，包含编码器和解码器两部分组成。AutoEncoder常被用于降维，关于降维可参考该网站。

AutoEncoder本质上还是一种线性变换，实现起来也比较简单。本文使用AutoEncoder来对MNIST数据集进行学习和生成。

代码如下：
首先是网络模型的定义，命名为auto_encoder.py

import torch
from torch import nn


class AutoEncoder(nn.Module):
    def __init__(self):
        super(AutoEncoder, self).__init__()

        # [b, 784]--->[b, 20]
        # 编码器 学习特征
        self.encoder = nn.Sequential(
            nn.Linear(784, 256),
            nn.ReLU(),
            nn.Linear(256, 64),
            nn.ReLU(),
            nn.Linear(64, 20),
            nn.ReLU()
        )

        # [b, 20]--->[b, 784]
        # 解码器 生成
        self.decoder = nn.Sequential(
            nn.Linear(20, 64),
            nn.ReLU(),
            nn.Linear(64, 256),
            nn.ReLU(),
            nn.Linear(256, 784),
            nn.Sigmoid()
        )

    def forward(self, x):
        # x: [b, 1, 28, 28]
        batch_size = x.size(0)
        # flatten
        x = x.view(batch_size, -1)
        # encoder
        x = self.encoder(x)
        # decoder
        x = self.decoder(x)
        # reshape
        x = x.view(batch_size, 1, 28, 28)

        return x

接下来写主函数main().py

import torch
from torch import nn, optim
from torch.utils.data import DataLoader
from torchvision import transforms, datasets

from auto_encoder import AutoEncoder  # 从上一个程序中导入网络模型

import visdom  # 使用visdom做可视化

def main():
	# 加载数据集
    mnist_train = datasets.MNIST(root='mnist', train=True, transform=transforms.Compose([
        transforms.ToTensor()]), download=True)
    mnist_train_loader = DataLoader(mnist_train, batch_size=32, shuffle=True)

    mnist_test = datasets.MNIST(root='mnist', train=True, transform=transforms.Compose([
        transforms.ToTensor()]), download=True)
    mnist_test_loader = DataLoader(mnist_test, batch_size=32, shuffle=False)

	
    x, _ = iter(mnist_train_loader).next()  # 因为是无监督学习 所以不使用label标签
    print("x.shape: ", x.shape)

    # 定义损失 选择优化器
    device = torch.device('cuda')
    model = AutoEncoder().to(device)  # 不要忘了转移到GPU上
    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=1e-3)
    print(model)

    viz = visdom.Visdom()  # 实例化visdom
	
	# train
    for epoch in range(50):
        for batch_idx, (x, _) in enumerate(mnist_train_loader):  
            # x: [b, 1, 28, 28]
            x = x.to(device)

            x_hat = model(x)
            loss = criterion(x_hat, x)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        print("Epoch: ", epoch, "Loss: ", loss.item())


        x, _ = iter(mnist_test_loader).next()  # 测试集上验证
        x = x.to(device)
        with torch.no_grad():
            x_hat = model(x)

        viz.images(x, nrow=8, win='x', opts=dict(title='x'))
        viz.images(x_hat, nrow=8, win='x_hat', opts=dict(title='x_hat'))
        # 每行画8个图

if __name__ == '__main__':
    main()

结果如下：