| | """ |
| | 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 |
