train.py

import numpy as np
from PIL import Image
from math import log, sqrt, pi
import argparse

import torch
from torch import nn, optim
from torch.autograd import Variable, grad
from torch.utils.data import DataLoader
from torchvision import datasets, transforms, utils
import wandb

from model import Glow

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def sample_data(path, batch_size, image_size):
    transform = transforms.Compose(
        [
            transforms.Resize(image_size),
            # transforms.CenterCrop(image_size),
            # transforms.RandomHorizontalFlip(),
            # transforms.Grayscale(num_output_channels=3),
            transforms.ToTensor(),
        ]
    )

    dataset = datasets.MNIST(root=path, train=True, transform=transform, download=True)
    loader = DataLoader(dataset, shuffle=True, batch_size=batch_size, num_workers=4)
    loader = iter(loader)

    while True:
        try:
            yield next(loader)

        except StopIteration:
            loader = DataLoader(
                dataset, shuffle=True, batch_size=batch_size, num_workers=4
            )
            loader = iter(loader)
            yield next(loader)


def calc_z_shapes(n_channel, input_size, n_flow, n_block):
    z_shapes = []

    for i in range(n_block - 1):
        input_size //= 2
        n_channel *= 2

        z_shapes.append((n_channel, input_size, input_size))

    input_size //= 2
    z_shapes.append((n_channel * 4, input_size, input_size))

    return z_shapes


def calc_loss(log_p, logdet, image_size, n_bins):
    # log_p = calc_log_p([z_list])
    # n_pixel = image_size * image_size * 3   # if img is RGB
    n_pixel = image_size * image_size

    loss = -log(n_bins) * n_pixel
    loss = loss + logdet + log_p

    return (
        (-loss / (log(2) * n_pixel)).mean(),
        (log_p / (log(2) * n_pixel)).mean(),
        (logdet / (log(2) * n_pixel)).mean(),
    )


def train(args, model, optimizer):
    dataset = iter(sample_data(args.path, args.batch, args.img_size))
    n_bins = 2.0 ** args.n_bits

    z_sample = []
    # z_shapes = calc_z_shapes(3, args.img_size, args.n_flow, args.n_block)   # if img is RGB
    z_shapes = calc_z_shapes(1, args.img_size, args.n_flow, args.n_block)
    for z in z_shapes:
        z_new = torch.randn(args.n_sample, *z) * args.temp
        z_sample.append(z_new.to(device))

    for i in range(args.iter):
        image, _ = next(dataset)
        image = image.to(device)

        image = image * 255

        if args.n_bits < 8:
            image = torch.floor(image / 2 ** (8 - args.n_bits))

        image = image / n_bins - 0.5

        if i == 0:
            with torch.no_grad():
                log_p, logdet, _ = model.module(
                    image + torch.rand_like(image) / n_bins
                )

                continue

        else:
            log_p, logdet, _ = model(image + torch.rand_like(image) / n_bins)

        logdet = logdet.mean()

        loss, log_p, log_det = calc_loss(log_p, logdet, args.img_size, n_bins)
        model.zero_grad()
        loss.backward()
        # warmup_lr = args.lr * min(1, i * batch_size / (50000 * 10))
        warmup_lr = args.lr
        optimizer.param_groups[0]["lr"] = warmup_lr
        optimizer.step()

        if i % 1 == 0:
            print(
                f"Loss: {loss.item():.5f}; logP: {log_p.item():.5f}; logdet: {log_det.item():.5f}; lr: {warmup_lr:.7f}"
            )
            wandb.log({"Loss": loss.item(), "logP": log_p.item(), "logdet": log_det.item(), "lr": warmup_lr})

        if i % 1000 == 0:
            with torch.no_grad():
                generate = model_single.reverse(z_sample).cpu().data
                utils.save_image(
                    generate,
                    f"sample/{str(i + 1).zfill(6)}.png",
                    normalize=True,
                    nrow=10,
                    range=(-0.5, 0.5),
                )
                wandb.log({
                    "images": wandb.Image(generate),  # 接收的是一个numpy格式的数组
                })

        if i % 1000 == 0:
            torch.save(
                model.state_dict(), f"checkpoint/model_{str(i + 1).zfill(6)}.pt"
            )
            torch.save(
                optimizer.state_dict(), f"checkpoint/optim_{str(i + 1).zfill(6)}.pt"
            )


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Glow trainer")
    parser.add_argument("--batch", default=32, type=int, help="batch size")
    parser.add_argument("--iter", default=100000, type=int, help="maximum iterations")
    parser.add_argument(
        "--n_flow", default=4, type=int, help="number of flows in each block"
    )
    parser.add_argument("--n_block", default=3, type=int, help="number of blocks")
    parser.add_argument(
        "--no_lu",
        action="store_true",
        help="use plain convolution instead of LU decomposed version",
    )
    parser.add_argument(
        "--affine", action="store_true", help="use affine coupling instead of additive"
    )
    parser.add_argument("--n_bits", default=5, type=int, help="number of bits")
    parser.add_argument("--lr", default=1e-4, type=float, help="learning rate")
    parser.add_argument("--img_size", default=56, type=int, help="image size")
    parser.add_argument("--temp", default=0.5, type=float, help="temperature of sampling")
    parser.add_argument("--n_sample", default=30, type=int, help="number of samples")
    parser.add_argument("--path", default="./mnist_data/", type=str, help="Path to image directory")

    args = parser.parse_args()
    print(args)

    # use wandb
    wandb.init(
        # set the wandb project where this run will be logged
        project="glow",

        # track hyperparameters and run metadata
        config={
            "architecture": "glow",
            "dataset": "MNIST",
            # hyperparameters
            "batch": args.batch,
            "iter": args.iter,
            "n_flow": args.n_flow,
            "n_block": args.n_block,
            "no_lu": args.no_lu,
            "affine": args.affine,
            "n_bits": args.n_bits,
            "lr": args.lr,
            "img_size": args.img_size,
            "temp": args.temp,
            "n_sample": args.n_sample,
        }
    )

    model_single = Glow(
        # 3, args.n_flow, args.n_block, affine=args.affine, conv_lu=not args.no_lu  # if img is RGB
        1, args.n_flow, args.n_block, affine=args.affine, conv_lu=not args.no_lu
    )
    model = nn.DataParallel(model_single)
    # model = model_single
    model = model.to(device)

    optimizer = optim.Adam(model.parameters(), lr=args.lr)

    train(args, model, optimizer)