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
9 # Any copyright is dedicated to the Public Domain.
10 # https://creativecommons.org/publicdomain/zero/1.0/
12 # Written by Francois Fleuret <francois@fleuret.org>
14 import sys, os, argparse, time, math, itertools
16 import torch, torchvision
18 from torch import optim, nn
19 from torch.nn import functional as F
21 ######################################################################
23 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
25 ######################################################################
27 parser = argparse.ArgumentParser(
28 description="Very simple implementation of a VAE for teaching."
31 parser.add_argument("--nb_epochs", type=int, default=25)
33 parser.add_argument("--learning_rate", type=float, default=1e-3)
35 parser.add_argument("--batch_size", type=int, default=100)
37 parser.add_argument("--data_dir", type=str, default="./data/")
39 parser.add_argument("--log_filename", type=str, default="train.log")
41 parser.add_argument("--latent_dim", type=int, default=32)
43 parser.add_argument("--nb_channels", type=int, default=128)
45 parser.add_argument("--no_dkl", action="store_true")
47 # With that option, do not follow the setup of the original VAE paper
48 # of forcing the variance of X|Z to 1 during training and to 0 for
49 # sampling, but optimize and use the variance.
50 parser.add_argument("--no_hacks", action="store_true")
52 args = parser.parse_args()
54 log_file = open(args.log_filename, "w")
56 ######################################################################
60 t = time.strftime("%Y-%m-%d_%H:%M:%S - ", time.localtime())
62 if log_file is not None:
63 log_file.write(t + s + "\n")
70 ######################################################################
73 def sample_gaussian(mu, log_var):
74 std = log_var.mul(0.5).exp()
75 return torch.randn(mu.size(), device=mu.device) * std + mu
78 def log_p_gaussian(x, mu, log_var):
81 (-0.5 * ((x - mu).pow(2) / var) - 0.5 * log_var - 0.5 * math.log(2 * math.pi))
87 def dkl_gaussians(mean_a, log_var_a, mean_b, log_var_b):
88 mean_a, log_var_a = mean_a.flatten(1), log_var_a.flatten(1)
89 mean_b, log_var_b = mean_b.flatten(1), log_var_b.flatten(1)
90 var_a = log_var_a.exp()
91 var_b = log_var_b.exp()
93 log_var_b - log_var_a - 1 + (mean_a - mean_b).pow(2) / var_b + var_a / var_b
97 ######################################################################
100 class LatentGivenImageNet(nn.Module):
101 def __init__(self, nb_channels, latent_dim):
104 self.model = nn.Sequential(
105 nn.Conv2d(1, nb_channels, kernel_size=1), # to 28x28
106 nn.ReLU(inplace=True),
107 nn.Conv2d(nb_channels, nb_channels, kernel_size=5), # to 24x24
108 nn.ReLU(inplace=True),
109 nn.Conv2d(nb_channels, nb_channels, kernel_size=5), # to 20x20
110 nn.ReLU(inplace=True),
111 nn.Conv2d(nb_channels, nb_channels, kernel_size=4, stride=2), # to 9x9
112 nn.ReLU(inplace=True),
113 nn.Conv2d(nb_channels, nb_channels, kernel_size=3, stride=2), # to 4x4
114 nn.ReLU(inplace=True),
115 nn.Conv2d(nb_channels, 2 * latent_dim, kernel_size=4),
118 def forward(self, x):
119 output = self.model(x).view(x.size(0), 2, -1)
120 mu, log_var = output[:, 0], output[:, 1]
124 class ImageGivenLatentNet(nn.Module):
125 def __init__(self, nb_channels, latent_dim):
128 self.model = nn.Sequential(
129 nn.ConvTranspose2d(latent_dim, nb_channels, kernel_size=4),
130 nn.ReLU(inplace=True),
132 nb_channels, nb_channels, kernel_size=3, stride=2
134 nn.ReLU(inplace=True),
136 nb_channels, nb_channels, kernel_size=4, stride=2
138 nn.ReLU(inplace=True),
139 nn.ConvTranspose2d(nb_channels, nb_channels, kernel_size=5), # from 20x20
140 nn.ReLU(inplace=True),
141 nn.ConvTranspose2d(nb_channels, 2, kernel_size=5), # from 24x24
144 def forward(self, z):
145 output = self.model(z.view(z.size(0), -1, 1, 1))
146 mu, log_var = output[:, 0:1], output[:, 1:2]
147 if not args.no_hacks:
152 ######################################################################
154 data_dir = os.path.join(args.data_dir, "mnist")
156 train_set = torchvision.datasets.MNIST(data_dir, train=True, download=True)
157 train_input = train_set.data.view(-1, 1, 28, 28).float()
159 test_set = torchvision.datasets.MNIST(data_dir, train=False, download=True)
160 test_input = test_set.data.view(-1, 1, 28, 28).float()
162 ######################################################################
164 model_q_Z_given_x = LatentGivenImageNet(
165 nb_channels=args.nb_channels, latent_dim=args.latent_dim
168 model_p_X_given_z = ImageGivenLatentNet(
169 nb_channels=args.nb_channels, latent_dim=args.latent_dim
172 optimizer = optim.Adam(
173 itertools.chain(model_p_X_given_z.parameters(), model_q_Z_given_x.parameters()),
174 lr=args.learning_rate,
177 model_p_X_given_z.to(device)
178 model_q_Z_given_x.to(device)
180 ######################################################################
182 train_input, test_input = train_input.to(device), test_input.to(device)
184 train_mu, train_std = train_input.mean(), train_input.std()
185 train_input.sub_(train_mu).div_(train_std)
186 test_input.sub_(train_mu).div_(train_std)
188 ######################################################################
190 mean_p_Z = train_input.new_zeros(1, args.latent_dim)
191 log_var_p_Z = mean_p_Z
193 for epoch in range(args.nb_epochs):
196 for x in train_input.split(args.batch_size):
197 mean_q_Z_given_x, log_var_q_Z_given_x = model_q_Z_given_x(x)
198 z = sample_gaussian(mean_q_Z_given_x, log_var_q_Z_given_x)
199 mean_p_X_given_z, log_var_p_X_given_z = model_p_X_given_z(z)
202 log_q_z_given_x = log_p_gaussian(z, mean_q_Z_given_x, log_var_q_Z_given_x)
203 log_p_x_z = log_p_gaussian(
204 x, mean_p_X_given_z, log_var_p_X_given_z
205 ) + log_p_gaussian(z, mean_p_Z, log_var_p_Z)
206 loss = -(log_p_x_z - log_q_z_given_x).mean()
208 log_p_x_given_z = log_p_gaussian(x, mean_p_X_given_z, log_var_p_X_given_z)
209 dkl_q_Z_given_x_from_p_Z = dkl_gaussians(
210 mean_q_Z_given_x, log_var_q_Z_given_x, mean_p_Z, log_var_p_Z
212 loss = (-log_p_x_given_z + dkl_q_Z_given_x_from_p_Z).mean()
214 optimizer.zero_grad()
218 acc_loss += loss.item() * x.size(0)
220 log_string(f"acc_loss {epoch} {acc_loss/train_input.size(0)}")
222 ######################################################################
225 def save_image(x, filename):
226 x = x * train_std + train_mu
227 x = x.clamp(min=0, max=255) / 255
228 torchvision.utils.save_image(1 - x, filename, nrow=16, pad_value=0.8)
231 # Save a bunch of test images
234 save_image(x, "input.png")
236 # Save the same images after encoding / decoding
238 mean_q_Z_given_x, log_var_q_Z_given_x = model_q_Z_given_x(x)
239 z = sample_gaussian(mean_q_Z_given_x, log_var_q_Z_given_x)
240 mean_p_X_given_z, log_var_p_X_given_z = model_p_X_given_z(z)
242 x = sample_gaussian(mean_p_X_given_z, log_var_p_X_given_z)
245 save_image(x, "output.png")
247 # Generate a bunch of images
249 z = sample_gaussian(mean_p_Z.expand(x.size(0), -1), log_var_p_Z.expand(x.size(0), -1))
250 mean_p_X_given_z, log_var_p_X_given_z = model_p_X_given_z(z)
252 x = sample_gaussian(mean_p_X_given_z, log_var_p_X_given_z)
255 save_image(x, "synth.png")
257 ######################################################################