pytorch（二十六）：自动编码器

一、自动编码器

1、AE.py

import torch
from torch import nn

class AE(nn.Module):
    def __init__(self):
        super(AE, 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):
        """
        :param x: [b, 1, 28, 28]
        :return:
        """
        batchsz = x.shape[0]
        #flatten
        x = x.view(batchsz, 784)
        #encoder
        x = self.encoder(x)
        #decoder
        x = self.decoder(x)
        #reshape
        x = x.view(batchsz,1, 28, 28)

        return x, None

2、main.py

import torch
from torch.utils.data import DataLoader
from torchvision import transforms, datasets
from auto_encoder import AE
from torch import nn, optim
import visdom
def main():
    mnist_train = datasets.MNIST("mnist", True, transform=transforms.Compose([
        transforms.ToTensor()
    ]), download=True)
    mnist_train = DataLoader(mnist_train, batch_size=32, shuffle=True)

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

    x, _ = iter(mnist_train).__next__()
    print(x.shape)

    model = AE()
    criton = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=1e-3)
    viz = visdom.Visdom()
    for epoch in range(1000):
        for batchidx, (x, _) in enumerate(mnist_train):
            #[b, 1, 28, 28]
            x_hat, _ = model(x)
            loss = criton(x_hat, x)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        print(epoch, "loss:",  loss.item())
        x, _ = iter(mnist_test).__next__()
        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"))

if __name__ == '__main__':
    main()

二、变分自动编码器编码器

1、模型

import torch
from torch import nn
import numpy as np

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

        #[b, 784] => [b, 20]
        #u:[b, 10]
        #sigma:[b, 10]
        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(10, 64),
            nn.ReLU(),
            nn.Linear(64, 256),
            nn.ReLU(),
            nn.Linear(256, 784),
            nn.Sigmoid(),
        )

    def forward(self, x):
        """
        :param x: [b, 1, 28, 28]
        :return:
        """
        batchsz = x.shape[0]
        #flatten
        x = x.view(batchsz, 784)
        #encoder
        #[b, 20], including mean and sigma
        h_ = self.encoder(x)
        #[b, 20] => [b, 10] and [b, 10]
        mu, sigma = h_.chunk(2, dim = 1)
        # reparametrize trick, epison~N(0, 1), [b, 10]
        h = mu + sigma * torch.randn_like(sigma)

        kld = 0.5 * torch.sum(
            torch.pow(mu, 2) +
            torch.pow(sigma, 2) -
            torch.log(1e-8 + torch.pow(sigma, 2)) - 1
        ) / (batchsz*28*28)

        #decoder
        x = self.decoder(h)
        #reshape
        x = x.view(batchsz,1, 28, 28)

        return x, kld

2、运行程序

import torch
from torch.utils.data import DataLoader
from torchvision import transforms, datasets
from vae import VAE
from torch import nn, optim
import visdom
def main():
    mnist_train = datasets.MNIST("mnist", True, transform=transforms.Compose([
        transforms.ToTensor()
    ]), download=True)
    mnist_train = DataLoader(mnist_train, batch_size=32, shuffle=True)

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

    x, _ = iter(mnist_train).__next__()
    print(x.shape)

    model = VAE()
    criton = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=1e-3)
    viz = visdom.Visdom()
    for epoch in range(1000):
        for batchidx, (x, _) in enumerate(mnist_train):
            #[b, 1, 28, 28]
            x_hat, kld = model(x)
            loss = criton(x_hat, x)

            if kld is not  None:
                loss = loss + 1.0 * kld

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        print(epoch, "loss:",  loss.item(), kld.item())
        x, _ = iter(mnist_test).__next__()
        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"))

if __name__ == '__main__':
    main()