diffvg/apps/sketch_gan.py

"""A simple training interface using ttools."""
import argparse
import os
import logging
import random

import numpy as np
import torch
from torchvision.datasets import MNIST
import torchvision.transforms as xforms
from torch.utils.data import DataLoader

import ttools
import ttools.interfaces

import pydiffvg

LOG = ttools.get_logger(__name__)

pydiffvg.render_pytorch.print_timing = False

torch.manual_seed(123)
np.random.seed(123)
torch.backends.cudnn.deterministic = True

latent_dim = 100
img_size = 32
num_paths = 8
num_segments = 8

def weights_init_normal(m):
    classname = m.__class__.__name__
    if classname.find("Conv") != -1:
        torch.nn.init.normal_(m.weight.data, 0.0, 0.02)
    elif classname.find("BatchNorm2d") != -1:
        torch.nn.init.normal_(m.weight.data, 1.0, 0.02)
        torch.nn.init.constant_(m.bias.data, 0.0)

class VisdomImageCallback(ttools.callbacks.ImageDisplayCallback):
    def visualized_image(self, batch, fwd_result):
        return torch.cat([batch[0], fwd_result.cpu()], dim = 2)

# From https://github.com/eriklindernoren/PyTorch-GAN/blob/master/implementations/dcgan/dcgan.py
class Generator(torch.nn.Module):
    def __init__(self):
        super(Generator, self).__init__()

        self.fc = torch.nn.Sequential(
            torch.nn.Linear(latent_dim, 128),
            torch.nn.LeakyReLU(0.2, inplace=True),
            torch.nn.Linear(128, 256),
            torch.nn.LeakyReLU(0.2, inplace=True),
            torch.nn.Linear(256, 512),
            torch.nn.LeakyReLU(0.2, inplace=True),
            torch.nn.Linear(512, 1024),
            torch.nn.LeakyReLU(0.2, inplace=True),
            torch.nn.Linear(1024, 2 * num_paths * (num_segments + 1) + num_paths + num_paths),
            torch.nn.Sigmoid()
        )

    def forward(self, z):
        out = self.fc(z)
        # construct paths
        imgs = []
        for b in range(out.shape[0]):
            index = 0
            shapes = []
            shape_groups = []
            for i in range(num_paths):
                points = img_size * out[b, index: index + 2 * (num_segments + 1)].view(-1, 2).cpu()
                index += 2 * (num_segments + 1)
                stroke_width = img_size * out[b, index].view(1).cpu()
                index += 1
                
                num_control_points = torch.zeros(num_segments, dtype = torch.int32) + 2
                path = pydiffvg.Path(num_control_points = num_control_points,
                                     points = points,
                                     stroke_width = stroke_width,
                                     is_closed = False)
                shapes.append(path)
    
                stroke_color = out[b, index].view(1).cpu()
                index += 1
                stroke_color = torch.cat([stroke_color, torch.tensor([0.0, 0.0, 1.0])])
                path_group = pydiffvg.ShapeGroup(shape_ids = torch.tensor([len(shapes) - 1]),
                                                 fill_color = None,
                                                 stroke_color = stroke_color)
                shape_groups.append(path_group)
            scene_args = pydiffvg.RenderFunction.serialize_scene(img_size, img_size, shapes, shape_groups)
            render = pydiffvg.RenderFunction.apply
            img = render(img_size, # width
                         img_size, # height
                         2,   # num_samples_x
                         2,   # num_samples_y
                         random.randint(0, 1048576), # seed
                         None,
                         *scene_args)
            img = img[:, :, :1]
            # HWC -> NCHW
            img = img.unsqueeze(0)
            img = img.permute(0, 3, 1, 2) # NHWC -> NCHW
            imgs.append(img)
        img = torch.cat(imgs, dim = 0)
        return img

class Discriminator(torch.nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()

        def discriminator_block(in_filters, out_filters, bn=True):
            block = [torch.nn.Conv2d(in_filters, out_filters, 3, 2, 1),
                     torch.nn.LeakyReLU(0.2, inplace=True),
                     torch.nn.Dropout2d(0.25)]
            if bn:
                block.append(torch.nn.BatchNorm2d(out_filters, 0.8))
            return block

        self.model = torch.nn.Sequential(
            *discriminator_block(1, 16, bn=False),
            *discriminator_block(16, 32),
            *discriminator_block(32, 64),
            *discriminator_block(64, 128),
        )

        # The height and width of downsampled image
        ds_size = img_size // 2 ** 4
        self.adv_layer = torch.nn.Sequential(
            torch.nn.Linear(128 * ds_size ** 2, 1),
            torch.nn.Sigmoid())

    def forward(self, img):
        out = self.model(img)
        out = out.view(out.shape[0], -1)
        validity = self.adv_layer(out)

        return validity

class MNISTInterface(ttools.interfaces.SGANInterface):
    """An adapter to run or train a model."""

    def __init__(self, gen, discrim, lr=2e-4):
        super(MNISTInterface, self).__init__(gen, discrim, lr, opt = 'adam')

    def forward(self, batch):
        return self.gen(torch.zeros([batch[0].shape[0], latent_dim], device = self.device).normal_())

    def _discriminator_input(self, batch, fwd_data, fake=False):
        if fake:
            return fwd_data
        else:
            return batch[0].to(self.device)

def train(args):
    """Train a MNIST classifier."""

    # Setup train and val data
    _xform = xforms.Compose([xforms.Resize([32, 32]), xforms.ToTensor()])
    data = MNIST("data/mnist", train=True, download=True, transform=_xform)

    # Initialize asynchronous dataloaders
    loader = DataLoader(data, batch_size=args.bs, num_workers=2)

    # Instantiate the models
    gen = Generator()
    discrim = Discriminator()

    gen.apply(weights_init_normal)
    discrim.apply(weights_init_normal)

    # Checkpointer to save/recall model parameters
    checkpointer_gen = ttools.Checkpointer(os.path.join(args.out, "checkpoints"), model=gen, prefix="gen_")
    checkpointer_discrim = ttools.Checkpointer(os.path.join(args.out, "checkpoints"), model=discrim, prefix="discrim_")

    # resume from a previous checkpoint, if any
    checkpointer_gen.load_latest()
    checkpointer_discrim.load_latest()

    # Setup a training interface for the model
    interface = MNISTInterface(gen, discrim, lr=args.lr)

    # Create a training looper with the interface we defined
    trainer = ttools.Trainer(interface)

    # Adds several callbacks, that will be called by the trainer --------------
    # A periodic checkpointing operation
    trainer.add_callback(ttools.callbacks.CheckpointingCallback(checkpointer_gen))
    trainer.add_callback(ttools.callbacks.CheckpointingCallback(checkpointer_discrim))
    # A simple progress bar
    trainer.add_callback(ttools.callbacks.ProgressBarCallback(
        keys=["loss_g", "loss_d", "loss"]))
    # A volatile logging using visdom
    trainer.add_callback(ttools.callbacks.VisdomLoggingCallback(
        keys=["loss_g", "loss_d", "loss"],
        port=8080, env="mnist_demo"))
    # Image
    trainer.add_callback(VisdomImageCallback(port=8080, env="mnist_demo"))
    # -------------------------------------------------------------------------

    # Start the training
    LOG.info("Training started, press Ctrl-C to interrupt.")
    trainer.train(loader, num_epochs=args.epochs)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    # TODO: subparsers
    parser.add_argument("data", help="directory where we download and store the MNIST dataset.")
    parser.add_argument("out", help="directory where we write the checkpoints and visualizations.")
    parser.add_argument("--lr", type=float, default=1e-4, help="learning rate for the optimizer.")
    parser.add_argument("--epochs", type=int, default=500, help="number of epochs to train for.")
    parser.add_argument("--bs", type=int, default=64, help="number of elements per batch.")
    args = parser.parse_args()
    ttools.set_logger(True)  # activate debug prints
    train(args)