class ConvAutoencoder(nn.Module):
def __init__(self, depth=4, dim_latent=64, dim_img=64, in_channels=3):
"""
depth: int
How many convolutional layers there are in the encoder/decoder
dim_latent: int
Dimension of the latent space
dim_digit: int
Width/height of input image
in_channels: int
Number of channels of input image
"""
super().__init__()
self.dim_latent = dim_latent
## Step 1: Create convolutional encoder
in_orig = in_channels
layers = []
out_channels = 64
for i in range(depth):
layers.append(nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=2, padding=1))
layers.append(nn.LeakyReLU())
in_channels = out_channels
out_channels *= 2
# Create a dummy input to get the shape right
X = torch.zeros(1, in_orig, dim_img, dim_img)
XOut = nn.Sequential(*layers)(X)
shape = XOut.shape[1::]
layers += [nn.Flatten()]
self.encoder = nn.Sequential(*layers)
# Create a linear layer for both mu and sigma
self.mu_enc = nn.Sequential(nn.Linear(np.prod(shape), dim_latent), nn.Tanh())
self.sigma_enc = nn.Sequential(nn.Linear(np.prod(shape), dim_latent), nn.ReLU())
## Step 2: Setup convolutional decoder
layers = [nn.Linear(dim_latent, np.prod(shape)), nn.LeakyReLU(), nn.Unflatten(1, shape)]
in_channels = out_channels//2
for i in range(depth):
out_channels = 3
if i < depth-1:
out_channels = in_channels // 2
# Use upsampling with bilinear interpolation instead of ConvTranspose
# to avoid checkerboard artifacts
# See this link for more info: https://distill.pub/2016/deconv-checkerboard/
layers.append(nn.Upsample(scale_factor=2, mode='bilinear'))
layers.append(nn.Conv2d(in_channels, out_channels, 3, stride=1, padding=1))
layers.append(nn.LeakyReLU())
in_channels = out_channels
self.decoder = nn.Sequential(*layers)
def forward(self, X):
flat_dim = self.encoder(X)
# Estimate mu and sigma
mu = 5*self.mu_enc(flat_dim)
sigma = self.sigma_enc(flat_dim)
# Use reparam trick to sample a z in the latent space
z = mu + sigma*torch.randn(sigma.shape).to(sigma)
XOut = self.decoder(z) # Decoding
# Fit the data
loss_fit = torch.sum((X-XOut)**2)
# Fit the prior (mu=0, sigma=1)
kl_div = 0.5*torch.sum(sigma**2 + mu**2 - 1 - torch.log(1e-8 + sigma**2), dim=1)
kl_div = torch.mean(kl_div)
return z, XOut, (loss_fit, kl_div)
def sample(self, n_examples, device):
"""
Sample from the latent space and generate the images from the decoder
"""
# Generate samples from Gaussian with mean 0, sigma 1
z = torch.randn(n_examples, self.dim_latent).to(device)
return self.decoder(z)
def plot_samples(XGen):
k = int(np.sqrt(XGen.shape[0]))
for i in range(XGen.shape[0]):
plt.subplot(k, k, i+1)
Xi = XGen[i, :, :, :].detach().cpu().numpy()
Xi = np.moveaxis(Xi, 0, 2)
Xi[Xi < 0] = 0
Xi[Xi > 1] = 1
plt.imshow(Xi)
plt.axis("off")
