自编码器生成MNIST(Pytorch)

发表于

自编码器介绍

自编码器主要是由编码器(Encoder)和解码器(Decoder)组成，它是一个试图去还原原始输入的一个系统。在深度学习中，自编码器是一种无监督的神经网络模型。它可以学习到输入数据的隐含特征，同时通过学习到的特征可以重构出原始数据。它类似于PCA，可以起到特征提取器的功能。

1

自编码器生成MNIST

导入相关模块
1
2
3
4
5
6
7
8
import torch
import torch.nn as nn
import torch.utils.data as data
import os
from torchvision import transforms
from torchvision.utils import save_image
from torchvision.datasets import MNIST
import torch.optim as optim
参数定义
1
2
3
4
5
6
epochs = 10
batch_size = 100
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
transforms = transforms.Compose([
        transforms.ToTensor(),
    ])
建立文件夹保存数据
1
2
3
4
5
6
if not os.path.exists("ae_img"):
    os.mkdir("ae_img")
if not os.path.exists("test_ae_img"):
    os.mkdir("test_ae_img")
if not os.path.exists("params"):
    os.mkdir("params")
加载数据集
1
2
3
4
mydataset = MNIST(root="./mnist_data/",train=True,download=True,transform=transforms)
dataloader = data.DataLoader(dataset=mydataset,shuffle=True,batch_size=batch_size)
test_data = MNIST(root="./mnist_data/", train=False, download=False, transform=transforms)
test_loader = data.DataLoader(dataset=test_data, shuffle=True, batch_size=batch_size)
定义编码器网络

在这里编码器网络采用两个卷积层和一个全连接层。并进行了归一化处理，激活函数使用ReLU激活函数。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class EncoderNet(nn.Module):
    def __init__(self):
        super(EncoderNet,self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(1,3,3,2,1),
            nn.BatchNorm2d(3),
            nn.ReLU(),
        )#N,3,14,14
        self.conv2 = nn.Sequential(
            nn.Conv2d(3,6,3,2,1),
            nn.BatchNorm2d(6),
            nn.ReLU(),
        )#N,6,7,7
        self.fc = nn.Sequential(
            nn.Linear(6*7*7,128),
        )#N,128

    def forward(self, x):
        y1 = self.conv1(x)
        y2 = self.conv2(y1)
        # print("encoder_y2",y2.shape) #[100, 6, 7, 7]
        y2 = torch.reshape(y2,[y2.size(0),-1]) #要经过全连接层，将4维转为2维,[100,294]
        out = self.fc(y2)
        return out
定义解码器网络

解码器网络和编码器网络相反，拿到编码器的输出，通过一个全连接层放大维度，之后是通过两个逆卷积层扩大图像的尺寸，扩大到28*28和MNIST原始数据集尺寸大小相同。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class DecoderNet(nn.Module):
    def __init__(self):
        super(DecoderNet,self).__init__()
        self.fc = nn.Sequential(
            nn.Linear(128,6*7*7),
            nn.BatchNorm1d(6*7*7),
            nn.ReLU()
        )#7,7
        self.conv1 = nn.Sequential(
            nn.ConvTranspose2d(6,3,3,2,1,output_padding=1),
            nn.BatchNorm2d(3),
            nn.ReLU()
        )#14,14
        self.conv2 = nn.Sequential(
            nn.ConvTranspose2d(3,1,3,2,1,output_padding=1),
            nn.ReLU()
        )#28,28
    def forward(self, x):
        y1 = self.fc(x) #[100,294]
        # print("decoder_y1", y1.shape)
        y1 = torch.reshape(y1,[y1.size(0),6,7,7]) #全连接之后，从二维转为四维
        y2 = self.conv1(y1)
        out = self.conv2(y2)
        return out
定义自编码器网络
1
2
3
4
5
6
7
8
9
class Net(nn.Module):
    def __init__(self):
        super(Net,self).__init__()
        self.encoder = EncoderNet()
        self.decoder = DecoderNet()
    def forward(self, x):
        encoder_out = self.encoder(x)
        decoder_out = self.decoder(encoder_out)
        return decoder_out
实例化模型
1
2
model = Net()
model = model.to(device)
定义优化器和损失函数

损失函数我们这里使用的是均方损失函数，因为这里我们衡量的是真实图片和自编码器生成的图片的损失，采用均方差计算损失会好一些。

1
2
optimizer = optim.Adam(model.parameters(),0.001)
criterion = nn.MSELoss()
定义训练函数

在每一轮训练中加载训练数据集，并计算真实数据和生成数据的损失，之后进行梯度优化。并将真实图片和生成的图片保存在文件夹中。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
def train_model(model, train_loader,optimizer ,num_epochs):
    for epoch in range(num_epochs):
        for i, (data, label) in enumerate(train_loader):
            img = data.to(device)
            out_img = model(img)
            loss = criterion(out_img, img)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            if i % 100 == 0:
                print("epoch:{},iteration:{}/{},loss:{:.3f}".format(epoch, i, len(dataloader), loss.float()))
            fake_image = out_img.cpu().data
            real_image = img.cpu().data
            save_image(fake_image, "./ae_img/epoch-{}-fake_img.jpg".format(epoch), nrow=10)
            save_image(real_image, "./ae_img/epoch-{}-real_img.jpg".format(epoch), nrow=10)
保存训练模型
1
torch.save(model.state_dict(), "./params/model.pth")
定义测试函数
1
2
3
4
5
6
7
8
9
10
11
12
13
14
def test_model(model,test_loader,num_epochs):
    model.eval()
    for epoch in range(num_epochs):
        for i, (img, label) in enumerate(test_loader):
            img = img.to(device)
            out_img = model(img)
            loss = criterion(out_img, img)

            if i % 100 == 0:
                print("epoch:{},loss:{:.3f}".format(epoch, loss.float()))
            fake_image = out_img.cpu().data
            real_image = img.cpu().data
            save_image(fake_image, "./test_ae_img/epoch-{}-fake_img.jpg".format(epoch), nrow=10)
            save_image(real_image, "./test_ae_img/epoch-{}-real_img.jpg".format(epoch), nrow=10)
模型测试
1
2
3
if __name__ == '__main__':
    train_model(model,dataloader,optimizer,epochs)
    test_model(model,test_loader,epochs)
数据展示

训练数据

2

3

测试数据

4

5