train_freyface.py

import os
import math
import torch
from torchvision import transforms
from torchvision.utils import save_image

from datasets import FreyFaceDataset
from models import VAE
from utils import make_gif, plot_elbocurve


# Random seed setting
torch.manual_seed(777)

# Device configuration
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Build the data input pipeline
entire_dataset = FreyFaceDataset(root='./data/FreyFace', transform=transforms.ToTensor())
batch_size = 100
indices = torch.randperm(len(entire_dataset)).tolist()
test_indices, train_indices = indices[:400], indices[400:]
train_dataset = torch.utils.data.Subset(dataset=entire_dataset, indices=train_indices)
test_dataset = torch.utils.data.Subset(dataset=entire_dataset, indices=test_indices)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)

# Build the model by instantiating the class "VAE"
input_size = 560
hidden_size = 200
latent_size = 2
model = VAE(input_size, hidden_size, latent_size, data_type='real').to(device)

pi = torch.Tensor([math.pi])
def compute_elbo(x, mean, log_var, out_mean, out_logvar):
    # ELBO(Evidence Lower Bound) is the objective of VAE, we train the model just to maximize the ELBO.
    
    reconst_error = -0.5 * torch.sum(torch.log(2*pi) + out_logvar + (x - out_mean).pow(2) / out_logvar.exp())
    # see Appendix B from VAE paper: "Kingma and Welling. Auto-Encoding Variational Bayes. ICLR-2014."
    # -KL[q(z|x)||p(z)] = 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
    kl_divergence = -0.5 * torch.sum(1 + log_var - mean.pow(2) - log_var.exp())
    elbo = (reconst_error - kl_divergence) / len(x)
    return elbo

# Select the optimizer
learning_rate = 1e-3
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

# Create a folder to store the experiment results if it doesn't exist
results_dir = 'results/FreyFace'
if not os.path.exists(results_dir):
    os.makedirs(results_dir)

# Save samples generated by the model before training
counter = 0
noise = torch.randn(35, latent_size).to(device)
generated_imgs = model.decode(noise)[0].view(-1, 1, 28, 20)
save_image(generated_imgs, os.path.join(results_dir, 'samples-0.png'), nrow=7)

# Start training
num_epochs = 2000
train_elbo = []
test_elbo = []
for epoch in range(1, num_epochs + 1):
    
    # model.train()   # optional, only useful when your model includes BatchNorm layers or Dropout layers
    for batch_idx, batch_x in enumerate(train_loader):
        # foward pass
        batch_x = batch_x.to(device).view(-1, input_size)
        z_mean, z_logvar, (output_mean, output_logvar) = model(batch_x)
        aver_loss = -compute_elbo(batch_x, z_mean, z_logvar, output_mean, output_logvar)
        # backprop and optimize
        optimizer.zero_grad()
        aver_loss.backward()
        optimizer.step()
        # print the average loss on batches 
        if (batch_idx + 1) % 4 == 0:
            print('Epoch {}/{}, Batch {}/{}, Aver_Loss: {:.2f}'.format(
                epoch, num_epochs, batch_idx + 1, math.ceil(len(train_dataset) / batch_size), aver_loss.item()))

    # model.eval()   # optional, corresponding to "model.train()"
    with torch.no_grad():
        # elbo_curve on training-set
        total_elbo = 0
        for batch_idx, batch_x in enumerate(train_loader):
            batch_x = batch_x.to(device).view(-1, input_size)
            z_mean, z_logvar, (output_mean, output_logvar) = model(batch_x)
            total_elbo += compute_elbo(batch_x, z_mean, z_logvar, output_mean, output_logvar).item()
        aver_elbo = total_elbo / (batch_idx + 1)
        train_elbo.append(aver_elbo)
        # elbo_curve on test-set 
        total_elbo = 0
        for batch_idx, batch_x in enumerate(test_loader):
            batch_x = batch_x.to(device).view(-1, input_size)
            z_mean, z_logvar, (output_mean, output_logvar) = model(batch_x)
            total_elbo += compute_elbo(batch_x, z_mean, z_logvar, output_mean, output_logvar).item()
        aver_elbo = total_elbo / (batch_idx + 1)
        test_elbo.append(aver_elbo)
        # save samples generated by the model at different stages of training
        if epoch == 1 or epoch == 2 or epoch == 4 or epoch == 8 or epoch == 16 or epoch == 32 or epoch == 64:
            counter += 1
            generated_imgs = model.decode(noise)[0].view(-1, 1, 28, 20)
            save_image(generated_imgs, os.path.join(results_dir, 'samples-{}.png'.format(counter)), nrow=7)
        if epoch % 100 == 0:
            counter += 1
            generated_imgs = model.decode(noise)[0].view(-1, 1, 28, 20)
            save_image(generated_imgs, os.path.join(results_dir, 'samples-{}.png'.format(counter)), nrow=7)


# Save the trained model's parameters
paras_dir = 'trained_parameters'
if not os.path.exists(paras_dir):
    os.makedirs(paras_dir)
torch.save(model.state_dict(), os.path.join(paras_dir, 'freyface_zdim{}.pkl'.format(latent_size)))

# Make a GIF using the samples generated by the model during training
make_gif(results_dir, counter + 1)

# Plot the elbo-curve on both the training set and the test set
plot_elbocurve(train_elbo, test_elbo, latent_size, results_dir)