Pytorch-自编码器与变分自编码器

提前导包：

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

import  visdom

1.自编码器（Auto-Encoder）

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):                 #x.shape=[b, 1, 28, 28]

        batchsz = x.size(0)
        x = x.view(batchsz, 784)          #flatten     
        x = self.encoder(x)               #encoder [b, 20]      
        x = self.decoder(x)               #decoder [b, 784]       
        x = x.view(batchsz, 1, 28, 28)    #reshape [b, 1, 28, 28]

        return x, None

2.变分自动编码器（Variational Auto-Encoder）

代码中的h和图中的c_i，计算方法略有不同，代码中没有用指数。

KL散度计算公式（代码中与sigma相乘的torch.randn_like(sigma)符合正态分布）：

class VAE(nn.Module):

    def __init__(self):
        super(VAE, 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(10, 64),
            nn.ReLU(),
            nn.Linear(64, 256),
            nn.ReLU(),
            nn.Linear(256, 784),
            nn.Sigmoid()
        )

        self.criteon = nn.MSELoss()

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

        kld = 0.5 * torch.sum(mu**2 + sigma**2 - torch.log(1e-8 + sigma**2) - 1) / (batchsz*28*28)   #KL散度计算
        
        return x_hat, kld

3.MINIST数据集上分别调用上面的编码器

def main():
    mnist_train = datasets.MNIST('mnist', train=True, transform=transforms.Compose([transforms.ToTensor()]), download=True)
    mnist_train = DataLoader(mnist_train, batch_size=32, shuffle=True)

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

    x, _ = iter(mnist_train).next()    #x: torch.Size([32, 1, 28, 28]) _: torch.Size([32])

    model = AE()
    # model = VAE()
    
    criteon = nn.MSELoss()             #均方损失
    optimizer = optim.Adam(model.parameters(), lr=1e-3)
    print(model)

    viz = visdom.Visdom()

    for epoch in range(20):

        for batchidx, (x, _) in enumerate(mnist_train):
            
            x_hat, kld = model(x)
            loss = criteon(x_hat, x)        #x_hat和x的shape=[b, 1, 28, 28]

            if kld is not None:
                elbo = - loss - 1.0 * kld   #elbo为证据下界
                loss = - elbo
            
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        print(epoch, 'loss:', loss.item())
        # print(epoch, 'loss:', loss.item(), 'kld:', kld.item())

        x, _ = iter(mnist_test).next()
        
        with torch.no_grad():
            x_hat, kld = 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()

当调用AE时，

0 loss: 0.02397083304822445
1 loss: 0.024659520015120506
2 loss: 0.020393237471580505
3 loss: 0.01954815723001957
4 loss: 0.01639191433787346
5 loss: 0.01630600169301033
6 loss: 0.017990168184041977
7 loss: 0.01680954359471798
8 loss: 0.015895305201411247
9 loss: 0.01704774796962738
10 loss: 0.013867242261767387
11 loss: 0.015675727277994156
12 loss: 0.015580415725708008
13 loss: 0.015662500634789467
14 loss: 0.014532235451042652
15 loss: 0.01624385453760624
16 loss: 0.014668326824903488
17 loss: 0.015973586589097977
18 loss: 0.0157624501734972
19 loss: 0.01488522719591856

当调用VAE时，

0 loss: 0.06747999787330627 kld: 0.017223423346877098
1 loss: 0.06267592310905457 kld: 0.01792667806148529
2 loss: 0.06116900593042374 kld: 0.01845495030283928
3 loss: 0.05097544193267822 kld: 0.0076100630685687065
4 loss: 0.05512534826993942 kld: 0.008729029446840286
5 loss: 0.04558167979121208 kld: 0.008567653596401215
6 loss: 0.04628278315067291 kld: 0.008163649588823318
7 loss: 0.05536432936787605 kld: 0.008285009302198887
8 loss: 0.048810530453920364 kld: 0.009821291081607342
9 loss: 0.046619318425655365 kld: 0.009058271534740925
10 loss: 0.04698382318019867 kld: 0.009476056322455406
11 loss: 0.048784226179122925 kld: 0.008850691840052605
12 loss: 0.05204786732792854 kld: 0.008851360529661179
13 loss: 0.04309754818677902 kld: 0.008809098042547703
14 loss: 0.05094045773148537 kld: 0.008593044243752956
15 loss: 0.04640775918960571 kld: 0.00919229444116354
16 loss: 0.04617678374052048 kld: 0.009322990663349628
17 loss: 0.044559232890605927 kld: 0.00912649929523468
18 loss: 0.04573676362633705 kld: 0.009612892754375935
19 loss: 0.040917910635471344 kld: 0.008869696408510208