#!/bin/env python
"""Train a Sketch-VAE."""
import argparse
from enum import Enum
import os
import wget
import time
import numpy as np
import torch as th
from torch.utils.data import DataLoader
import torchvision.datasets as dset
import torchvision.transforms as transforms
import ttools
import ttools.interfaces
from ttools.modules import networks
import rendering
import losses
import modules
import data
import pydiffvg
LOG = ttools.get_logger(__name__)
BASE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), os.pardir)
OUTPUT = os.path.join(BASE_DIR, "results")
class SketchVAE(th.nn.Module):
class ImageEncoder(th.nn.Module):
def __init__(self, image_size=64, width=64, zdim=128):
super(SketchVAE.ImageEncoder, self).__init__()
self.zdim = zdim
self.net = th.nn.Sequential(
th.nn.Conv2d(4, width, 5, padding=2),
th.nn.InstanceNorm2d(width),
th.nn.ReLU(inplace=True),
# 64x64
th.nn.Conv2d(width, width, 5, padding=2),
th.nn.InstanceNorm2d(width),
th.nn.ReLU( inplace=True),
# 64x64
th.nn.Conv2d(width, 2*width, 5, stride=1, padding=2),
th.nn.InstanceNorm2d(2*width),
th.nn.ReLU( inplace=True),
# 32x32
th.nn.Conv2d(2*width, 2*width, 5, stride=2, padding=2),
th.nn.InstanceNorm2d(2*width),
th.nn.ReLU( inplace=True),
# 16x16
th.nn.Conv2d(2*width, 2*width, 5, stride=2, padding=2),
th.nn.InstanceNorm2d(2*width),
th.nn.ReLU( inplace=True),
# 16x16
th.nn.Conv2d(2*width, 2*width, 5, stride=2, padding=2),
th.nn.InstanceNorm2d(2*width),
th.nn.ReLU( inplace=True),
# 8x8
th.nn.Conv2d(2*width, 2*width, 5, stride=2, padding=2),
th.nn.InstanceNorm2d(2*width),
th.nn.ReLU( inplace=True),
# 4x4
modules.Flatten(),
th.nn.Linear(4*4*2*width, 2*zdim)
)
def forward(self, images):
features = self.net(images)
# VAE params
mu = features[:, :self.zdim]
log_sigma = features[:, self.zdim:]
# Sample a latent vector
sigma = th.exp(log_sigma/2.0)
z0 = th.randn(self.zdim, device=mu.device)
z = mu + sigma*z0
# KL divergence needs mu/sigma
return z, mu, log_sigma
class ImageDecoder(th.nn.Module):
""""""
def __init__(self, zdim=128, image_size=64, width=64):
super(SketchVAE.ImageDecoder, self).__init__()
self.zdim = zdim
self.width = width
self.embedding = th.nn.Linear(zdim, 4*4*2*width)
self.net = th.nn.Sequential(
th.nn.ConvTranspose2d(2*width, 2*width, 4, padding=1, stride=2),
th.nn.InstanceNorm2d(2*width),
th.nn.ReLU( inplace=True),
# 8x8
th.nn.ConvTranspose2d(2*width, 2*width, 4, padding=1, stride=2),
th.nn.InstanceNorm2d(2*width),
th.nn.ReLU( inplace=True),
# 16x16
th.nn.ConvTranspose2d(2*width, 2*width, 4, padding=1, stride=2),
th.nn.InstanceNorm2d(2*width),
th.nn.ReLU( inplace=True),
# 16x16
th.nn.Conv2d(2*width, 2*width, 5, padding=2, stride=1),
th.nn.InstanceNorm2d(2*width),
th.nn.ReLU( inplace=True),
# 16x16
th.nn.ConvTranspose2d(2*width, 2*width, 4, padding=1, stride=2),
th.nn.InstanceNorm2d(2*width),
th.nn.ReLU( inplace=True),
# 32x32
th.nn.Conv2d(2*width, width, 5, padding=2, stride=1),
th.nn.InstanceNorm2d(width),
th.nn.ReLU( inplace=True),
# 32x32
th.nn.ConvTranspose2d(width, width, 5, padding=2, stride=1),
th.nn.InstanceNorm2d(width),
th.nn.ReLU( inplace=True),
# 64x64
th.nn.Conv2d(width, width, 5, padding=2, stride=1),
th.nn.InstanceNorm2d(width),
th.nn.ReLU( inplace=True),
# 64x64
th.nn.Conv2d(width, 4, 5, padding=2, stride=1),
)
def forward(self, z):
bs = z.shape[0]
im = self.embedding(z).view(bs, 2*self.width, 4, 4)
out = self.net(im)
return out
class SketchDecoder(th.nn.Module):
"""
The decoder outputs a sequence where each time step models (dx, dy,
opacity).
"""
def __init__(self, sequence_length, hidden_size=512, dropout=0.9,
zdim=128, num_layers=3):
super(SketchVAE.SketchDecoder, self).__init__()
self.sequence_length = sequence_length
self.hidden_size = hidden_size
self.num_layers = num_layers
self.zdim = zdim
# Maps the latent vector to an initial cell/hidden vector
self.hidden_cell_predictor = th.nn.Linear(zdim, 2*hidden_size*num_layers)
self.lstm = th.nn.LSTM(
zdim, hidden_size,
num_layers=self.num_layers, dropout=dropout,
batch_first=True)
self.dxdy_predictor = th.nn.Sequential(
th.nn.Linear(hidden_size, 2),
th.nn.Tanh(),
)
self.opacity_predictor = th.nn.Sequential(
th.nn.Linear(hidden_size, 1),
th.nn.Sigmoid(),
)
def forward(self, z, hidden_and_cell=None):
# Every step in the sequence takes the latent vector as input so we
# replicate it here
bs = z.shape[0]
steps = self.sequence_length - 1 # no need to predict the start of sequence
expanded_z = z.unsqueeze(1).repeat(1, steps, 1)
if hidden_and_cell is None:
# Initialize from latent vector
hidden_and_cell = self.hidden_cell_predictor(
th.tanh(z))
hidden = hidden_and_cell[:, :self.hidden_size*self.num_layers]
hidden = hidden.view(-1, self.num_layers, self.hidden_size)
hidden = hidden.permute(1, 0, 2).contiguous()
# hidden = hidden.unsqueeze(1).contiguous()
cell = hidden_and_cell[:, self.hidden_size*self.num_layers:]
cell = cell.view(-1, self.num_layers, self.hidden_size)
cell = cell.permute(1, 0, 2).contiguous()
# cell = cell.unsqueeze(1).contiguous()
hidden_and_cell = (hidden, cell)
outputs, hidden_and_cell = self.lstm(expanded_z, hidden_and_cell)
hidden, cell = hidden_and_cell
dxdy = self.dxdy_predictor(
outputs.reshape(bs*steps, self.hidden_size)).view(bs, steps, -1)
opacity = self.opacity_predictor(
outputs.reshape(bs*steps, self.hidden_size)).view(bs, steps, -1)
strokes = th.cat([dxdy, opacity], -1)
return strokes
def __init__(self, sequence_length, zdim=128, image_size=64):
super(SketchVAE, self).__init__()
self.im_encoder = SketchVAE.ImageEncoder(
zdim=zdim, image_size=image_size)
self.im_decoder = SketchVAE.ImageDecoder(
zdim=zdim, image_size=image_size)
self.sketch_decoder = SketchVAE.SketchDecoder(
sequence_length, zdim=zdim)
def forward(self, images):
# Encode the images as latent vectors
z, mu, log_sigma = self.im_encoder(images)
decoded_im = self.im_decoder(z)
decoded_sketch = self.sketch_decoder(z)
return {
"decoded_im": decoded_im,
"decoded_sketch": decoded_sketch,
"z": z,
"mu": mu,
"log_sigma": log_sigma,
}
class SketchVAECallback(ttools.callbacks.ImageDisplayCallback):
"""Simple callback that visualize images."""
def visualized_image(self, batch, step_data, is_val=False):
if is_val:
return None
# only display the first n drawings
n = 8
gt = step_data["gt_image"][:n].detach()
vae_im = step_data["vae_image"][:n].detach()
sketch_im = step_data["sketch_image"][:n].detach()
rendering = th.cat([gt, vae_im, sketch_im], 2)
rendering = th.clamp(rendering, 0, 1)
alpha = rendering[:, 3:4]
rendering = rendering[:, :3] * alpha
return rendering
def caption(self, batch, step_data, is_val=False):
if is_val:
return ""
else:
return "top: truth, middle: vae sample, output: rnn-output"
class Interface(ttools.ModelInterface):
def __init__(self, model, lr=1e-4, lr_decay=0.9999,
kl_weight=0.5, kl_min_weight=0.01, kl_decay=0.99995,
raster_resolution=64, absolute_coords=False,
device="cpu", grad_clip=1.0):
super(Interface, self).__init__()
self.grad_clip = grad_clip
self.raster_resolution = raster_resolution
self.absolute_coords = absolute_coords
self.model = model
self.device = device
self.model.to(self.device)
self.im_enc_opt = th.optim.Adam(
self.model.im_encoder.parameters(), lr=lr)
self.im_dec_opt = th.optim.Adam(
self.model.im_decoder.parameters(), lr=lr)
self.sketch_dec_opt = th.optim.Adam(
self.model.sketch_decoder.parameters(), lr=lr)
self.kl_weight = kl_weight
self.kl_min_weight = kl_min_weight
self.kl_decay = kl_decay
self.kl_loss = losses.KLDivergence()
self.schedulers = [
th.optim.lr_scheduler.ExponentialLR(self.im_enc_opt, lr_decay),
th.optim.lr_scheduler.ExponentialLR(self.im_dec_opt, lr_decay),
th.optim.lr_scheduler.ExponentialLR(self.sketch_dec_opt, lr_decay),
]
# include loss on alpha
self.im_loss = losses.MultiscaleMSELoss(channels=4).to(self.device)
def optimizers(self):
return [self.im_enc_opt, self.im_dec_opt, self.sketch_dec_opt]
def kl_scaling(self):
# Scale the KL divergence weight
try:
state = self.im_enc_opt.state_dict()["param_groups"][0]["params"][0]
optim_step = self.im_enc_opt.state_dict()["state"][state]["step"]
except KeyError:
optim_step = 0 # no step taken yet
kl_scaling = 1.0 - (1.0 -
self.kl_min_weight)*(self.kl_decay**optim_step)
return kl_scaling
def training_step(self, batch):
gt_strokes, gt_im = batch
gt_strokes = gt_strokes.to(self.device)
gt_im = gt_im.to(self.device)
out = self.model(gt_im)
kl_loss = self.kl_loss(
out["mu"], out["log_sigma"])
kl_weight = self.kl_weight * self.kl_scaling()
# add start of sequence
sos = gt_strokes[:, :1]
sketch = th.cat([sos, out["decoded_sketch"]], 1)
vae_im = out["decoded_im"]
# start = time.time()
sketch_im = rendering.opacityStroke2diffvg(
sketch, canvas_size=self.raster_resolution, debug=False,
force_cpu=True, relative=not self.absolute_coords)
# elapsed = (time.time() - start)*1000
# print("out rendering took %.2fms" % elapsed)
vae_im_loss = self.im_loss(vae_im, gt_im)
sketch_im_loss = self.im_loss(sketch_im, gt_im)
# vae_im_loss = th.nn.functional.mse_loss(vae_im, gt_im)
# sketch_im_loss = th.nn.functional.mse_loss(sketch_im, gt_im)
loss = vae_im_loss + kl_loss*kl_weight + sketch_im_loss
self.im_enc_opt.zero_grad()
self.im_dec_opt.zero_grad()
self.sketch_dec_opt.zero_grad()
loss.backward()
# clip gradients
enc_nrm = th.nn.utils.clip_grad_norm_(
self.model.im_encoder.parameters(), self.grad_clip)
dec_nrm = th.nn.utils.clip_grad_norm_(
self.model.im_decoder.parameters(), self.grad_clip)
sketch_dec_nrm = th.nn.utils.clip_grad_norm_(
self.model.sketch_decoder.parameters(), self.grad_clip)
if enc_nrm > self.grad_clip:
LOG.debug("Clipped encoder gradient (%.5f) to %.2f",
enc_nrm, self.grad_clip)
if dec_nrm > self.grad_clip:
LOG.debug("Clipped decoder gradient (%.5f) to %.2f",
dec_nrm, self.grad_clip)
if sketch_dec_nrm > self.grad_clip:
LOG.debug("Clipped sketch decoder gradient (%.5f) to %.2f",
sketch_dec_nrm, self.grad_clip)
self.im_enc_opt.step()
self.im_dec_opt.step()
self.sketch_dec_opt.step()
return {
"vae_image": vae_im,
"sketch_image": sketch_im,
"gt_image": gt_im,
"loss": loss.item(),
"vae_im_loss": vae_im_loss.item(),
"sketch_im_loss": sketch_im_loss.item(),
"kl_loss": kl_loss.item(),
"kl_weight": kl_weight,
"lr": self.im_enc_opt.param_groups[0]["lr"],
}
def init_validation(self):
return dict(sample=None)
def validation_step(self, batch, running_data):
# Switch to eval mode for dropout, batchnorm, etc
# self.model.eval()
# with th.no_grad():
# # sample = self.model.sample(
# # batch.to(self.device), temperature=self.sampling_temperature)
# # running_data["sample"] = sample
# self.model.train()
return running_data
def train(args):
th.manual_seed(0)
np.random.seed(0)
dataset = data.FixedLengthQuickDrawDataset(
args.dataset, max_seq_length=args.sequence_length,
canvas_size=args.raster_resolution)
dataloader = DataLoader(
dataset, batch_size=args.bs, num_workers=args.workers, shuffle=True)
# val_dataset = [s for idx, s in enumerate(dataset) if idx < 8]
# val_dataloader = DataLoader(
# val_dataset, batch_size=8, num_workers=4, shuffle=False)
val_dataloader = None
model_params = {
"zdim": args.zdim,
"sequence_length": args.sequence_length,
"image_size": args.raster_resolution,
# "encoder_dim": args.encoder_dim,
# "decoder_dim": args.decoder_dim,
}
model = SketchVAE(**model_params)
model.train()
LOG.info("Model parameters:\n%s", model_params)
device = "cpu"
if th.cuda.is_available():
device = "cuda"
LOG.info("Using CUDA")
interface = Interface(model, raster_resolution=args.raster_resolution,
lr=args.lr, lr_decay=args.lr_decay,
kl_decay=args.kl_decay, kl_weight=args.kl_weight,
absolute_coords=args.absolute_coordinates,
device=device)
env_name = "sketch_vae"
if args.custom_name is not None:
env_name += "_" + args.custom_name
if args.absolute_coordinates:
env_name += "_abs_coords"
chkpt = os.path.join(OUTPUT, env_name)
# Resume from checkpoint, if any
checkpointer = ttools.Checkpointer(
chkpt, model, meta=model_params,
optimizers=interface.optimizers(),
schedulers=interface.schedulers)
extras, meta = checkpointer.load_latest()
epoch = extras["epoch"] if extras and "epoch" in extras.keys() else 0
if meta is not None and meta != model_params:
LOG.info("Checkpoint's metaparams differ "
"from CLI, aborting: %s and %s", meta, model_params)
trainer = ttools.Trainer(interface)
# Add callbacks
losses = ["loss", "kl_loss", "vae_im_loss", "sketch_im_loss"]
training_debug = ["lr", "kl_weight"]
trainer.add_callback(ttools.callbacks.ProgressBarCallback(
keys=losses, val_keys=None))
trainer.add_callback(ttools.callbacks.VisdomLoggingCallback(
keys=losses, val_keys=None, env=env_name, port=args.port))
trainer.add_callback(ttools.callbacks.VisdomLoggingCallback(
keys=training_debug, smoothing=0, val_keys=None, env=env_name,
port=args.port))
trainer.add_callback(ttools.callbacks.CheckpointingCallback(
checkpointer, max_files=2, interval=600, max_epochs=10))
trainer.add_callback(
ttools.callbacks.LRSchedulerCallback(interface.schedulers))
trainer.add_callback(SketchVAECallback(
env=env_name, win="samples", port=args.port, frequency=args.freq))
# Start training
trainer.train(dataloader, starting_epoch=epoch,
val_dataloader=val_dataloader,
num_epochs=args.num_epochs)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--dataset", default="cat.npz")
parser.add_argument("--absolute_coordinates", action="store_true",
default=False)
parser.add_argument("--custom_name")
# Training params
parser.add_argument("--bs", type=int, default=1)
parser.add_argument("--workers", type=int, default=0)
parser.add_argument("--num_epochs", type=int, default=10000)
parser.add_argument("--lr", type=float, default=1e-4)
parser.add_argument("--lr_decay", type=float, default=0.9999)
parser.add_argument("--kl_weight", type=float, default=0.5)
parser.add_argument("--kl_decay", type=float, default=0.99995)
# Model configuration
parser.add_argument("--zdim", type=int, default=128)
parser.add_argument("--sequence_length", type=int, default=50)
parser.add_argument("--raster_resolution", type=int, default=64)
# parser.add_argument("--encoder_dim", type=int, default=256)
# parser.add_argument("--decoder_dim", type=int, default=512)
# Viz params
parser.add_argument("--freq", type=int, default=10)
parser.add_argument("--port", type=int, default=5000)
args = parser.parse_args()
pydiffvg.set_use_gpu(False)
train(args)