Spaces:
Running
Running
File size: 3,067 Bytes
c53ddec |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
"""
Module containing the encoders.
"""
import numpy as np
import torch
from torch import nn
# ALL encoders should be called Enccoder<Model>
def get_encoder(model_type):
model_type = model_type.lower().capitalize()
return eval("Encoder{}".format(model_type))
class EncoderBurgess(nn.Module):
def __init__(self, img_size,
latent_dim=10):
r"""Encoder of the model proposed in [1].
Parameters
----------
img_size : tuple of ints
Size of images. E.g. (1, 32, 32) or (3, 64, 64).
latent_dim : int
Dimensionality of latent output.
Model Architecture (transposed for decoder)
------------
- 4 convolutional layers (each with 32 channels), (4 x 4 kernel), (stride of 2)
- 2 fully connected layers (each of 256 units)
- Latent distribution:
- 1 fully connected layer of 20 units (log variance and mean for 10 Gaussians)
References:
[1] Burgess, Christopher P., et al. "Understanding disentangling in
$\beta$-VAE." arXiv preprint arXiv:1804.03599 (2018).
"""
super(EncoderBurgess, self).__init__()
# Layer parameters
hid_channels = 32
kernel_size = 4
hidden_dim = 256
self.latent_dim = latent_dim
self.img_size = img_size
# Shape required to start transpose convs
self.reshape = (hid_channels, kernel_size, kernel_size)
n_chan = self.img_size[0]
# Convolutional layers
cnn_kwargs = dict(stride=2, padding=1)
self.conv1 = nn.Conv2d(n_chan, hid_channels, kernel_size, **cnn_kwargs)
self.conv2 = nn.Conv2d(hid_channels, hid_channels, kernel_size, **cnn_kwargs)
self.conv3 = nn.Conv2d(hid_channels, hid_channels, kernel_size, **cnn_kwargs)
# If input image is 64x64 do fourth convolution
if self.img_size[1] == self.img_size[2] == 64:
self.conv_64 = nn.Conv2d(hid_channels, hid_channels, kernel_size, **cnn_kwargs)
# Fully connected layers
self.lin1 = nn.Linear(np.product(self.reshape), hidden_dim)
self.lin2 = nn.Linear(hidden_dim, hidden_dim)
# Fully connected layers for mean and variance
self.mu_logvar_gen = nn.Linear(hidden_dim, self.latent_dim * 2)
def forward(self, x):
batch_size = x.size(0)
# Convolutional layers with ReLu activations
x = torch.relu(self.conv1(x))
x = torch.relu(self.conv2(x))
x = torch.relu(self.conv3(x))
if self.img_size[1] == self.img_size[2] == 64:
x = torch.relu(self.conv_64(x))
# Fully connected layers with ReLu activations
x = x.view((batch_size, -1))
x = torch.relu(self.lin1(x))
x = torch.relu(self.lin2(x))
# Fully connected layer for log variance and mean
# Log std-dev in paper (bear in mind)
mu_logvar = self.mu_logvar_gen(x)
mu, logvar = mu_logvar.view(-1, self.latent_dim, 2).unbind(-1)
return mu, logvar
|