tiny_vae.py

   1 #!/usr/bin/env python
   2
   3 # @XREMOTE_HOST: elk.fleuret.org
   4 # @XREMOTE_EXEC: python
   5 # @XREMOTE_PRE: source ${HOME}/misc/venv/pytorch/bin/activate
   6 # @XREMOTE_PRE: ln -sf ${HOME}/data/pytorch ./data
   7 # @XREMOTE_GET: *.png
   8
   9 # Any copyright is dedicated to the Public Domain.
  10 # https://creativecommons.org/publicdomain/zero/1.0/
  11
  12 # Written by Francois Fleuret <francois@fleuret.org>
  13
  14 import sys, os, argparse, time, math
  15
  16 import torch, torchvision
  17
  18 from torch import optim, nn
  19 from torch.nn import functional as F
  20
  21 ######################################################################
  22
  23 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  24
  25 ######################################################################
  26
  27 parser = argparse.ArgumentParser(
  28     description="Very simple implementation of a VAE for teaching."
  29 )
  30
  31 parser.add_argument("--nb_epochs", type=int, default=25)
  32
  33 parser.add_argument("--learning_rate", type=float, default=1e-3)
  34
  35 parser.add_argument("--batch_size", type=int, default=100)
  36
  37 parser.add_argument("--data_dir", type=str, default="./data/")
  38
  39 parser.add_argument("--log_filename", type=str, default="train.log")
  40
  41 parser.add_argument("--latent_dim", type=int, default=32)
  42
  43 parser.add_argument("--nb_channels", type=int, default=32)
  44
  45 parser.add_argument("--no_dkl", action="store_true")
  46
  47 parser.add_argument("--beta", type=float, default=1.0)
  48
  49 args = parser.parse_args()
  50
  51 log_file = open(args.log_filename, "w")
  52
  53 ######################################################################
  54
  55
  56 def log_string(s):
  57     t = time.strftime("%Y-%m-%d_%H:%M:%S ", time.localtime())
  58
  59     if log_file is not None:
  60         log_file.write(t + s + "\n")
  61         log_file.flush()
  62
  63     print(t + s)
  64     sys.stdout.flush()
  65
  66
  67 ######################################################################
  68
  69
  70 def sample_categorical(param):
  71     dist = torch.distributions.Categorical(logits=param)
  72     return (dist.sample().unsqueeze(1).float() - train_mu) / train_std
  73
  74
  75 def log_p_categorical(x, param):
  76     x = (x.squeeze(1) * train_std + train_mu).long().clamp(min=0, max=255)
  77     param = param.permute(0, 3, 1, 2)
  78     return -F.cross_entropy(param, x, reduction="none").flatten(1).sum(dim=1)
  79
  80
  81 def sample_gaussian(param):
  82     mean, log_var = param
  83     std = log_var.mul(0.5).exp()
  84     return torch.randn(mean.size(), device=mean.device) * std + mean
  85
  86
  87 def log_p_gaussian(x, param):
  88     mean, log_var, x = param[0].flatten(1), param[1].flatten(1), x.flatten(1)
  89     var = log_var.exp()
  90     return -0.5 * (((x - mean).pow(2) / var) + log_var + math.log(2 * math.pi)).sum(1)
  91
  92
  93 def dkl_gaussians(param_a, param_b):
  94     mean_a, log_var_a = param_a[0].flatten(1), param_a[1].flatten(1)
  95     mean_b, log_var_b = param_b[0].flatten(1), param_b[1].flatten(1)
  96     var_a = log_var_a.exp()
  97     var_b = log_var_b.exp()
  98     return 0.5 * (
  99         log_var_b - log_var_a - 1 + (mean_a - mean_b).pow(2) / var_b + var_a / var_b
 100     ).sum(1)
 101
 102
 103 def dup_param(param, nb):
 104     mean, log_var = param
 105     s = (nb,) + (-1,) * (mean.dim() - 1)
 106     return (mean.expand(s), log_var.expand(s))
 107
 108
 109 ######################################################################
 110
 111
 112 class VariationalAutoEncoder(nn.Module):
 113     def __init__(self, nb_channels, latent_dim):
 114         super().__init__()
 115
 116         self.encoder = nn.Sequential(
 117             nn.Conv2d(1, nb_channels, kernel_size=1),  # to 28x28
 118             nn.ReLU(inplace=True),
 119             nn.Conv2d(nb_channels, nb_channels, kernel_size=5),  # to 24x24
 120             nn.ReLU(inplace=True),
 121             nn.Conv2d(nb_channels, nb_channels, kernel_size=5),  # to 20x20
 122             nn.ReLU(inplace=True),
 123             nn.Conv2d(nb_channels, nb_channels, kernel_size=4, stride=2),  # to 9x9
 124             nn.ReLU(inplace=True),
 125             nn.Conv2d(nb_channels, nb_channels, kernel_size=3, stride=2),  # to 4x4
 126             nn.ReLU(inplace=True),
 127             nn.Conv2d(nb_channels, 2 * latent_dim, kernel_size=4),
 128         )
 129
 130         self.decoder = nn.Sequential(
 131             nn.ConvTranspose2d(latent_dim, nb_channels, kernel_size=4),
 132             nn.ReLU(inplace=True),
 133             nn.ConvTranspose2d(
 134                 nb_channels, nb_channels, kernel_size=3, stride=2
 135             ),  # from 4x4
 136             nn.ReLU(inplace=True),
 137             nn.ConvTranspose2d(
 138                 nb_channels, nb_channels, kernel_size=4, stride=2
 139             ),  # from 9x9
 140             nn.ReLU(inplace=True),
 141             nn.ConvTranspose2d(nb_channels, nb_channels, kernel_size=5),  # from 20x20
 142             nn.ReLU(inplace=True),
 143             nn.ConvTranspose2d(nb_channels, 2, kernel_size=5),  # from 24x24
 144         )
 145
 146     def encode(self, x):
 147         output = self.encoder(x).view(x.size(0), 2, -1)
 148         mu, log_var = output[:, 0], output[:, 1]
 149         return mu, log_var
 150
 151     def decode(self, z):
 152         # return self.decoder(z.view(z.size(0), -1, 1, 1)).permute(0, 2, 3, 1)
 153         output = self.decoder(z.view(z.size(0), -1, 1, 1))
 154         mu, log_var = output[:, 0:1], output[:, 1:2]
 155         log_var.flatten(1)[...] = 1  # math.log(1e-1)
 156         # log_var.flatten(1)[...] = log_var.flatten(1)[:, :1]
 157         # log_var = log_var.clamp(min=2*math.log(1/256))
 158         return mu, log_var
 159
 160
 161 ######################################################################
 162
 163 data_dir = os.path.join(args.data_dir, "mnist")
 164
 165 train_set = torchvision.datasets.MNIST(data_dir, train=True, download=True)
 166 train_input = train_set.data.view(-1, 1, 28, 28).float()
 167
 168 test_set = torchvision.datasets.MNIST(data_dir, train=False, download=True)
 169 test_input = test_set.data.view(-1, 1, 28, 28).float()
 170
 171 ######################################################################
 172
 173
 174 def save_images(model, prefix=""):
 175     def save_image(x, filename):
 176         x = x * train_std + train_mu
 177         x = x.clamp(min=0, max=255) / 255
 178         torchvision.utils.save_image(1 - x, filename, nrow=12, pad_value=1.0)
 179         log_string(f"wrote {filename}")
 180
 181     # Save a bunch of train images
 182
 183     x = train_input[:36]
 184     save_image(x, f"{prefix}train_input.png")
 185
 186     # Save the same images after encoding / decoding
 187
 188     param_q_Z_given_x = model.encode(x)
 189     z = sample_gaussian(param_q_Z_given_x)
 190     param_p_X_given_z = model.decode(z)
 191     x = sample_gaussian(param_p_X_given_z)
 192     save_image(x, f"{prefix}train_output.png")
 193     save_image(param_p_X_given_z[0], f"{prefix}train_output_mean.png")
 194
 195     # Save a bunch of test images
 196
 197     x = test_input[:36]
 198     save_image(x, f"{prefix}input.png")
 199
 200     # Save the same images after encoding / decoding
 201
 202     param_q_Z_given_x = model.encode(x)
 203     z = sample_gaussian(param_q_Z_given_x)
 204     param_p_X_given_z = model.decode(z)
 205     x = sample_gaussian(param_p_X_given_z)
 206     save_image(x, f"{prefix}output.png")
 207     save_image(param_p_X_given_z[0], f"{prefix}output_mean.png")
 208
 209     # Generate a bunch of images
 210
 211     z = sample_gaussian(dup_param(param_p_Z, x.size(0)))
 212     param_p_X_given_z = model.decode(z)
 213     x = sample_gaussian(param_p_X_given_z)
 214     save_image(x, f"{prefix}synth.png")
 215     save_image(param_p_X_given_z[0], f"{prefix}synth_mean.png")
 216
 217
 218 ######################################################################
 219
 220 model = VariationalAutoEncoder(nb_channels=args.nb_channels, latent_dim=args.latent_dim)
 221
 222 model.to(device)
 223
 224 ######################################################################
 225
 226 train_input, test_input = train_input.to(device), test_input.to(device)
 227
 228 train_mu, train_std = train_input.mean(), train_input.std()
 229 train_input.sub_(train_mu).div_(train_std)
 230 test_input.sub_(train_mu).div_(train_std)
 231
 232 ######################################################################
 233
 234 zeros = train_input.new_zeros(1, args.latent_dim)
 235
 236 param_p_Z = zeros, zeros
 237
 238 for n_epoch in range(args.nb_epochs):
 239     optimizer = optim.Adam(
 240         model.parameters(),
 241         lr=args.learning_rate,
 242     )
 243
 244     acc_loss = 0
 245
 246     for x in train_input.split(args.batch_size):
 247         param_q_Z_given_x = model.encode(x)
 248         z = sample_gaussian(param_q_Z_given_x)
 249         param_p_X_given_z = model.decode(z)
 250         log_p_x_given_z = log_p_gaussian(x, param_p_X_given_z)
 251
 252         if args.no_dkl:
 253             log_q_z_given_x = log_p_gaussian(z, param_q_Z_given_x)
 254             log_p_z = log_p_gaussian(z, param_p_Z)
 255             log_p_x_z = log_p_x_given_z + log_p_z
 256             loss = -(log_p_x_z - log_q_z_given_x).mean()
 257         else:
 258             dkl_q_Z_given_x_from_p_Z = dkl_gaussians(param_q_Z_given_x, param_p_Z)
 259             loss = -(log_p_x_given_z - args.beta * dkl_q_Z_given_x_from_p_Z).mean()
 260
 261         optimizer.zero_grad()
 262         loss.backward()
 263         optimizer.step()
 264
 265         acc_loss += loss.item() * x.size(0)
 266
 267     log_string(f"acc_loss {n_epoch} {acc_loss/train_input.size(0)}")
 268
 269     if (n_epoch + 1) % 25 == 0:
 270         save_images(model, f"epoch_{n_epoch+1:04d}_")
 271
 272 ######################################################################