Update.

vae.py

@@ -5,7 +5,6 @@ import torch
 import torch.nn as nn
 import torchvision.datasets
 import torch.nn.functional as F
-from torchvision import transforms
 
 
 class Encoder(nn.Module):
@@ -65,65 +64,3 @@ class VAE(nn.Module):
         c = mu + sigma * torch.randn_like(sigma)
         xhat = self.decoder(c)
         return xhat, mu, sigma
-
-
-if __name__ == '__main__':
-    N_EPOCHS = 50
-    LEARNING_RATE = .001
-
-    model = VAE().cuda()
-    optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)
-    loss_fn = torch.nn.MSELoss()
-
-    dataset_base = torchvision.datasets.FashionMNIST("MNIST", download=True, transform=transforms.ToTensor())
-    dataset_train, dataset_test = torch.utils.data.random_split(
-        dataset_base, (int(.8 * len(dataset_base)), int(.2 * len(dataset_base)))
-    )
-
-    model.train()
-    dataloader = torch.utils.data.DataLoader(dataset_train,
-            batch_size=512,
-            shuffle=True,
-            num_workers=0)
-    i = 0
-    for epoch in range(N_EPOCHS):
-        total_loss = 0
-        for x, label in dataloader:
-            x = x.cuda()
-            label = label.cuda()
-            optimizer.zero_grad()
-            xhat, mu, logvar = model(x)
-
-            BCE = F.binary_cross_entropy(xhat, x.reshape(xhat.shape), reduction='mean')
-            # https://arxiv.org/abs/1312.6114
-            # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
-            KLD = -0.5 * torch.mean(1 + logvar - mu.pow(2) - logvar.exp())
-            loss = BCE + KLD
-            loss.backward()
-            optimizer.step()
-            total_loss += loss.item()
-        print(f"{epoch}: {total_loss:.4f}")
-
-    model.cpu()
-    with open("vae.pt", "wb") as file:
-        torch.save(model, file)
-    model.eval()
-    dataloader = torch.utils.data.DataLoader(dataset_test,
-            batch_size=512,
-            shuffle=True,
-            num_workers=0)
-    n_correct = 0
-
-    COLS = 4
-    ROWS = 4
-    fig, axes = plt.subplots(ncols=COLS, nrows=ROWS, figsize=(5.5, 3.5),
-                        constrained_layout=True)
-
-    dataloader_gen = iter(dataloader)
-    x, label = next(dataloader_gen)
-    xhat, mu, logvar = model(x)
-    xhat = xhat.reshape((-1, 28, 28))
-    for row in range(ROWS):
-        for col in range(COLS):
-            axes[row, col].imshow(xhat[row * COLS + col].detach().numpy(), cmap="gray")
-    plt.show()