vae_train.py

import os
import warnings


warnings.filterwarnings("ignore")

import torch
from torch.utils.tensorboard import SummaryWriter
from torch.utils.data import DataLoader
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
import random
from torch.optim.lr_scheduler import ReduceLROnPlateau

from utils import LinearScheduler, save_args, vae_loss
from vae_model import GeneralVAE
from torch_dataloader import PerspectiveTransformTorchDataset
from utils import parse_args
from utils import init_weights

args = parse_args()


def evaluate(model, loader, kl_scheduler, writer, i_epoch, name, device):
    model.eval()
    all_loss = []
    all_recon_loss = []
    all_kl_loss = []
    for i_batch, (batch) in enumerate(loader):
        recon_loss, kl_loss, source, target, predict, mu, logsigma, mu_attn_map, logsigma_attn_map = model(batch)
        loss = recon_loss + kl_loss * args.beta * kl_scheduler.val()
        all_loss.append(loss.detach().cpu())
        all_recon_loss.append(recon_loss.detach().cpu())
        all_kl_loss.append(kl_loss.detach().cpu())

    mean_loss = torch.stack(all_loss).mean()
    mean_recon_loss = torch.stack(all_recon_loss).mean()
    mean_kl_loss = torch.stack(all_kl_loss).mean()

    print("Epoch %i %s loss: %.4f, recon_loss: %.4f, kl_loss: %.4f" % (
    i_epoch, name, mean_loss, mean_recon_loss, mean_kl_loss))
    writer.add_scalar('%s/loss' % name, mean_loss.detach().cpu().item(), i_epoch)
    writer.add_scalar('%s/recon_loss' % name, mean_recon_loss.detach().cpu().item(), i_epoch)
    writer.add_scalar('%s/kl_loss' % name, mean_kl_loss.detach().cpu().item(), i_epoch)

    if i_epoch % args.log_freq == 0:
        writer.add_images('%s/images/source' % name, source[:8], i_epoch)
        if args.seq and not args.stacked_frames:
            writer.add_images('%s/images/target' % name, target[:8], i_epoch)
            writer.add_images('%s/images/transform' % name, predict[:8], i_epoch)
        else:
            writer.add_images('%s/images/target' % name, target[:8], i_epoch)
            writer.add_images('%s/images/transform' % name, predict[:8], i_epoch)
    return mean_loss, mean_recon_loss, mean_kl_loss, source, source, predict


def main():
    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
    print("Running %s latent %i, with seed %i on device: %r" % (args.architecture, args.latent_size, args.seed, device))
    random.seed(args.seed)
    os.environ['PYTHONHASHSEED'] = str(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)
    torch.backends.cudnn.deterministic = True

    ########### Preparation ##########
    # 1. dataset and dataloader
    train_dataset = PerspectiveTransformTorchDataset(args, split="train")
    valid_dataset = PerspectiveTransformTorchDataset(args, split="valid")
    test_dataset = PerspectiveTransformTorchDataset(args, split="test")

    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.workers,
                              prefetch_factor=16 if args.workers else None, drop_last=True, pin_memory=True)
    valid_loader = DataLoader(valid_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.workers,
                              prefetch_factor=16 if args.workers else None, drop_last=True, pin_memory=True)
    test_loader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.workers,
                             prefetch_factor=16 if args.workers else None, drop_last=True, pin_memory=True)

    source_shape, target_shape = test_dataset.get_shape()
    # 2. model and optimizer
    model = GeneralVAE(source_shape, target_shape, device, latent_size=args.latent_size,
                       architecture=args.architecture, disable_loss_weight=args.disable_loss_weight,
                       attn=args.attn, use_meta=args.use_meta, seq=args.seq, stacked_frames=args.stacked_frames,
                       seq_len=args.seq_len).to(device)
    model.apply(init_weights)
    optim = torch.optim.Adam(model.parameters(), lr=args.lr)
    # scheduler = ReduceLROnPlateau(optim, 'min', patience=10, verbose=True)

    if args.load_path is not None:
        model.load_state_dict(torch.load(args.load_path, map_location=device))
        print("Loaded model from %s" % args.load_path)

    ########### Train ############
    # 3. log
    writer = SummaryWriter("%s/models/%s/%s/latent_%i/%s_seed_%i_background_%s" % (args.log_dir, args.seq_arch, args.architecture, args.latent_size, args.tag, args.seed, args.background))
    save_args(args, writer.log_dir)

    # 4. main train
    step = 0
    kl_scheduler = LinearScheduler(args.start_time, args.start_value, args.end_time, args.end_value)
    best_valid_recon_loss = float("Inf")
    patient = 0

    for i_epoch in range(args.epochs):
        print("Starting Epoch %i" % i_epoch)
        model.train()
        kl_scheduler.step()
        for i_batch, batch in enumerate(train_loader):
            recon_loss, kl_loss, source, target, predict, mu, logsigma, mu_attn_map, logsigma_attn_map = model(batch)
            loss = recon_loss + kl_loss * args.beta * kl_scheduler.val()
            optim.zero_grad()
            loss.backward()
            optim.step()
            step += 1

            if args.debug and i_batch > 2:
                break

        writer.add_scalar('train/loss', loss.detach().cpu().item(), i_epoch)
        writer.add_scalar('train/recon_loss', recon_loss.detach().cpu().item(), i_epoch)
        writer.add_scalar('train/kl_loss', kl_loss.detach().cpu().item(), i_epoch)
        writer.add_scalar('train/kl_annealing', kl_scheduler.val(), i_epoch)

        if i_epoch % args.log_freq == 0:
            writer.add_images('train/images/source', source[:8], i_epoch)
            if args.seq and not args.stacked_frames:
                writer.add_images('train/images/target', target[:8:args.seq_len], i_epoch)
                writer.add_images('train/images/transform', predict[:8:args.seq_len], i_epoch)
            else:
                writer.add_images('train/images/target', target[:8], i_epoch)
                writer.add_images('train/images/transform', predict[:8], i_epoch)

        valid_loss, valid_recon_loss, valid_kl_loss, _, _, _ = evaluate(model, valid_loader, kl_scheduler, writer, i_epoch, "valid", device)
        # scheduler.step(valid_recon_loss)

        # Test set
        test_loss, test_recon_loss, test_kl_loss, test_source, test_source, test_predict = evaluate(model, test_loader, kl_scheduler, writer, i_epoch, "test", device)
        print("test loss: %.4f, test recon_loss: %.4f, test kl_loss: %.4f" % (
            test_loss, test_recon_loss, test_kl_loss))

        if valid_recon_loss < best_valid_recon_loss:
            patient = 0
            best_valid_recon_loss = valid_recon_loss
            torch.save(model.state_dict(), '%s/model.pt' % writer.get_logdir())
            print("Saving model, recon_loss: %.4f" % valid_recon_loss)
        else:
            patient += 1
            if patient > args.es_patience:
                print("Valid loss increases, early stopping.")
                break

    # Test set
    test_loss, test_recon_loss, test_kl_loss, test_source, test_source, test_predict = evaluate(model, test_loader, kl_scheduler, writer, 0, "final test", device)
    print("Final test loss: %.4f, test recon_loss: %.4f, test kl_loss: %.4f" % (
        test_loss, test_recon_loss, test_kl_loss))

    writer.close()


if __name__ == '__main__':
    main()