1. import  torch
    2. from    torch import nn
    8. class AE(nn.Module):
    12.     def __init__(self):
    13.         super(AE, self).__init__()
    16.         # [b, 784] => [b, 20]
    17.         self.encoder = nn.Sequential(
    18.             nn.Linear(784, 256),
    19.             nn.ReLU(),
    20.             nn.Linear(256, 64),
    21.             nn.ReLU(),
    22.             nn.Linear(64, 20),
    23.             nn.ReLU()
    24.         )
    25.         # [b, 20] => [b, 784]
    26.         self.decoder = nn.Sequential(
    27.             nn.Linear(20, 64),
    28.             nn.ReLU(),
    29.             nn.Linear(64, 256),
    30.             nn.ReLU(),
    31.             nn.Linear(256, 784),
    32.             nn.Sigmoid()
    33.         )
    36.     def forward(self, x):
    37.         """
    39.         :param x: [b, 1, 28, 28]
    40.         :return:
    41.         """
    42.         batchsz = x.size(0)
    43.         # flatten
    44.         x = x.view(batchsz, 784)
    45.         # encoder
    46.         x = self.encoder(x)
    47.         # decoder
    48.         x = self.decoder(x)
    49.         # reshape
    50.         x = x.view(batchsz, 1, 28, 28)
    52.         return x, None
    1. import  torch
    2. from    torch.utils.data import DataLoader
    3. from    torch import nn, optim
    4. from    torchvision import transforms, datasets
    6. from    ae import AE
    7. from    vae import VAE
    9. import  visdom
    11. def main():
    12.     mnist_train = datasets.MNIST('mnist', True, transform=transforms.Compose([
    13.         transforms.ToTensor()
    14.     ]), download=True)
    15.     mnist_train = DataLoader(mnist_train, batch_size=32, shuffle=True)
    18.     mnist_test = datasets.MNIST('mnist', False, transform=transforms.Compose([
    19.         transforms.ToTensor()
    20.     ]), download=True)
    21.     mnist_test = DataLoader(mnist_test, batch_size=32, shuffle=True)
    24.     x, _ = iter(mnist_train).next()
    25.     print('x:', x.shape)
    27.     device = torch.device('cuda')
    28.     # model = AE().to(device)
    29.     model = VAE().to(device)
    30.     criteon = nn.MSELoss()
    31.     optimizer = optim.Adam(model.parameters(), lr=1e-3)
    32.     print(model)
    34.     viz = visdom.Visdom()
    36.     for epoch in range(1000):
    39.         for batchidx, (x, _) in enumerate(mnist_train):
    40.             # [b, 1, 28, 28]
    41.             x = x.to(device)
    43.             x_hat, kld = model(x)
    44.             loss = criteon(x_hat, x)
    46.             if kld is not None:
    47.                 elbo = - loss - 1.0 * kld
    48.                 loss = - elbo
    50.             # backprop
    51.             optimizer.zero_grad()
    52.             loss.backward()
    53.             optimizer.step()
    56.         print(epoch, 'loss:', loss.item(), 'kld:', kld.item())
    58.         x, _ = iter(mnist_test).next()
    59.         x = x.to(device)
    60.         with torch.no_grad():
    61.             x_hat, kld = model(x)
    62.         viz.images(x, nrow=8, win='x', opts=dict(title='x'))
    63.         viz.images(x_hat, nrow=8, win='x_hat', opts=dict(title='x_hat'))
    70. if __name__ == '__main__':
    71.     main()
    1. import  torch
    2. from    torch import nn
    8. class VAE(nn.Module):
    12.     def __init__(self):
    13.         super(VAE, self).__init__()
    16.         # [b, 784] => [b, 20]
    17.         # u: [b, 10]
    18.         # sigma: [b, 10]
    19.         self.encoder = nn.Sequential(
    20.             nn.Linear(784, 256),
    21.             nn.ReLU(),
    22.             nn.Linear(256, 64),
    23.             nn.ReLU(),
    24.             nn.Linear(64, 20),
    25.             nn.ReLU()
    26.         )
    27.         # [b, 20] => [b, 784]
    28.         self.decoder = nn.Sequential(
    29.             nn.Linear(10, 64),
    30.             nn.ReLU(),
    31.             nn.Linear(64, 256),
    32.             nn.ReLU(),
    33.             nn.Linear(256, 784),
    34.             nn.Sigmoid()
    35.         )
    37.         self.criteon = nn.MSELoss()
    39.     def forward(self, x):
    40.         """
    42.         :param x: [b, 1, 28, 28]
    43.         :return:
    44.         """
    45.         batchsz = x.size(0)
    46.         # flatten
    47.         x = x.view(batchsz, 784)
    48.         # encoder
    49.         # [b, 20], including mean and sigma
    50.         h_ = self.encoder(x)
    51.         # [b, 20] => [b, 10] and [b, 10]
    52.         mu, sigma = h_.chunk(2, dim=1)
    53.         # reparametrize trick, epison~N(0, 1)
    54.         h = mu + sigma * torch.randn_like(sigma)
    56.         # decoder
    57.         x_hat = self.decoder(h)
    58.         # reshape
    59.         x_hat = x_hat.view(batchsz, 1, 28, 28)
    61.         kld = 0.5 * torch.sum(
    62.             torch.pow(mu, 2) +
    63.             torch.pow(sigma, 2) -
    64.             torch.log(1e-8 + torch.pow(sigma, 2)) - 1
    65.         ) / (batchsz*28*28)
    67.         return x_hat, kld