变分自编码器实现（VAE）

import torch
from torch import nn

class VAE(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(784, 256),
            nn.ReLU(),
            nn.Linear(256, 64),
            nn.ReLU(),
            nn.Linear(64, 20),
            nn.ReLU()
        )

        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):
        batchsize = x.size(0)
        x = x.view(batchsize, 784)
        # [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
        h = mu + sigma * torch.randn_like(sigma)

        x_hat = self.decoder(h)
        x_hat = x_hat.view(batchsize, 1, 28, 28)

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

        return x_hat, kld

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

def main():
    vis_dom = visdom.Visdom()

    mnist_train = datasets.MNIST(root='minst', train=True, transform=transforms.Compose([
        transforms.ToTensor()
    ]), download=True, )

    mnist_train = DataLoader(mnist_train, batch_size=32, shuffle=True)

    mnist_test = datasets.MNIST(root='minst', train=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:', x.shape)

    model = VAE()
    criteon = nn.MSELoss()
    optimizer = optim.Adam(params=model.parameters(), lr=1e-3)

    for epoch in range(1000):
        epoch_loss = 0
        for batchidx, (x, _) in enumerate(mnist_train):
            x_hat, kld = model(x)
            loss = criteon(x_hat, x)
            if kld is not None:
                elbo = -loss - 1.0 * kld
                loss = -elbo
            epoch_loss = loss
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        print(epoch, 'loss:', epoch_loss, "kld",kld.item())

        x, _ = iter(mnist_test).next()
        with torch.no_grad():
            x_hat, kld = model(x)
        vis_dom.images(x, nrow=8, win='x', opts=dict(title='x'))
        vis_dom.images(x_hat, nrow=8, win='x_hat', opts=dict(title='x_hat'))


if __name__ == '__main__':
    main()

笔记参考课程地址：https://www.bilibili.com/video/BV1mT4y1F7u8