|
""" |
|
Full definition of a VAE model, all of it in this single file. |
|
References: |
|
1) An Introduction to Variational Autoencoders: |
|
https://arxiv.org/abs/1906.02691 |
|
""" |
|
|
|
|
|
import torch |
|
import torch.nn as nn |
|
|
|
|
|
class VAE(nn.Module): |
|
"""VAE for 64x64 face generation. |
|
|
|
The hidden dimensions can be tuned. |
|
""" |
|
|
|
def __init__(self, hiddens=[16, 32, 64, 128, 256], latent_dim=128) -> None: |
|
super().__init__() |
|
|
|
|
|
prev_channels = 3 |
|
modules = [] |
|
img_length = 64 |
|
for cur_channels in hiddens: |
|
modules.append( |
|
nn.Sequential( |
|
nn.Conv2d(prev_channels, |
|
cur_channels, |
|
kernel_size=3, |
|
stride=2, |
|
padding=1), nn.BatchNorm2d(cur_channels), |
|
nn.ReLU())) |
|
prev_channels = cur_channels |
|
img_length //= 2 |
|
self.encoder = nn.Sequential(*modules) |
|
self.mean_linear = nn.Linear(prev_channels * img_length * img_length, |
|
latent_dim) |
|
self.var_linear = nn.Linear(prev_channels * img_length * img_length, |
|
latent_dim) |
|
self.latent_dim = latent_dim |
|
|
|
|
|
modules = [] |
|
self.decoder_projection = nn.Linear( |
|
latent_dim, prev_channels * img_length * img_length) |
|
self.decoder_input_chw = (prev_channels, img_length, img_length) |
|
for i in range(len(hiddens) - 1, 0, -1): |
|
modules.append( |
|
nn.Sequential( |
|
nn.ConvTranspose2d(hiddens[i], |
|
hiddens[i - 1], |
|
kernel_size=3, |
|
stride=2, |
|
padding=1, |
|
output_padding=1), |
|
nn.BatchNorm2d(hiddens[i - 1]), nn.ReLU())) |
|
modules.append( |
|
nn.Sequential( |
|
nn.ConvTranspose2d(hiddens[0], |
|
hiddens[0], |
|
kernel_size=3, |
|
stride=2, |
|
padding=1, |
|
output_padding=1), |
|
nn.BatchNorm2d(hiddens[0]), nn.ReLU(), |
|
nn.Conv2d(hiddens[0], 3, kernel_size=3, stride=1, padding=1), |
|
nn.ReLU())) |
|
self.decoder = nn.Sequential(*modules) |
|
|
|
def forward(self, x): |
|
encoded = self.encoder(x) |
|
encoded = torch.flatten(encoded, 1) |
|
mean = self.mean_linear(encoded) |
|
logvar = self.var_linear(encoded) |
|
eps = torch.randn_like(logvar) |
|
std = torch.exp(logvar / 2) |
|
z = eps * std + mean |
|
x = self.decoder_projection(z) |
|
x = torch.reshape(x, (-1, *self.decoder_input_chw)) |
|
decoded = self.decoder(x) |
|
|
|
return decoded, mean, logvar |
|
|
|
def sample(self, device='cuda'): |
|
z = torch.randn(1, self.latent_dim).to(device) |
|
x = self.decoder_projection(z) |
|
x = torch.reshape(x, (-1, *self.decoder_input_chw)) |
|
decoded = self.decoder(x) |
|
return decoded |
