I am trying to implement a variational autoencoder, but calculating the Kullback Leibler divergence doesn’t work out the way I hoped.
Background: The input is a 1x800 tensor, which will be mapped to a 1x1200 tensor.
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from torch.utils.data import DataLoader, TensorDataset, IterableDataset
import os
from torch import optim
from model import Network
class Network(nn.Module):
def __init__(self, input_dim, output_dim, latent_dim, layer_dim1, layer_dim2):
"""
Parameter:
input_dim (int): number of inputs
output_dim (int): number of outputs
latent_dim (int): number of latent neurons
Layer_dim (int): number of neurons in hidden layers
"""
super(Network, self).__init__()
self.latent_dim = latent_dim
self.enc1 = nn.Linear(input_dim, layer_dim1)
self.enc2 = nn.Linear(layer_dim1, layer_dim2)
self.latent = nn.Linear(layer_dim2, latent_dim*2)
self.dec1 = nn.Linear(latent_dim, layer_dim2)
self.dec2 = nn.Linear(layer_dim2, layer_dim1)
self.out = nn.Linear(layer_dim1, output_dim)
def encoder(self, x):
z = F.elu(self.enc1(x))
z = F.elu(self.enc2(x))
z = self.latent(z)
self.mu = z[0:self.latent_dim]
self.log_sigma = z[self.latent_dim:]
self.sigma = torch.exp(self.log_sigma)
eps = torch.randn(x.size(0), self.latent_dim)
z_sample = self.mu + self.sigma*eps
self.kl_loss = kl_divergence(self.mu, self.log_sigma, dim=self.latent_dim)
return z_sample
def decoder(self, z):
x = F.elu(self.dec1(z))
x = F.elu(self.dec2(x))
return self.out(x)
def forward(self, batch):
self.latent_rep = self.encoder(batch)
dec_input = self.latent_rep
return self.decoder(dec_input)
def kl_divergence(means, log_sigma, dim, target_sigma=0.1):
"""
Computes Kullback–Leibler divergence for arrays of mean and log(sigma)
"""
target_sigma = torch.Tensor([target_sigma])
out = 1 / 2. * torch.mean(torch.mean(1 / target_sigma**2 * means**2 + torch.exp(2 * log_sigma) / target_sigma**2 - 2 * log_sigma + 2 * torch.log(target_sigma), dim=1) - dim)
out = out
return out
model = Network(800,1200,3,800,200)
SAVE_PATH = "trained/model.dat"
epochs = 5
learning_rate = 0.001
optimizer = optim.Adam(model.parameters(),lr=learning_rate, eps=1e-08)
hist_error = []
hist_loss = []
beta = 0.5
for epoch in range(epochs):
epoch_error = []
epoch_loss = []
for i in x:
optimizer.zero_grad()
print(i)
#i = torch.tensor(i)
#i.unsqueeze_(0)
pred = model.forward(i)
loss = torch.mean(torch.sum((pred - y[i]) ** 2))
loss.backward()
optimizer.step()
error = torch.mean(torch.sqrt((pred - y[i]) ** 2)).detach().numpy()
epoch_error.append(error)
epoch_loss.append(loss.data.detach().numpy())
hist_error.append(np.mean(epoch_error))
hist_loss.append(np.mean(epoch_loss))
print("Epoch %d -- loss %f, RMS error %f " % (epoch+1, hist_loss[-1], hist_error[-1]))
The error message I now get is something I looked into last night, but could not resolve:
Traceback (most recent call last):
File "/home/samim/miniconda3/envs/deep/lib/python3.6/site-packages/IPython/core/interactiveshell.py", line 3343, in run_code
exec(code_obj, self.user_global_ns, self.user_ns)
File "<ipython-input-2-54864ad18480>", line 1, in <module>
runfile('/home/samim/Documents/train.py', wdir='/home/samim/Documents')
File "/home/samim/.local/share/JetBrains/PyCharm2020.3/python/helpers/pydev/_pydev_bundle/pydev_umd.py", line 197, in runfile
pydev_imports.execfile(filename, global_vars, local_vars) # execute the script
File "/home/samim/.local/share/JetBrains/PyCharm2020.3/python/helpers/pydev/_pydev_imps/_pydev_execfile.py", line 18, in execfile
exec(compile(contents+"\n", file, 'exec'), glob, loc)
File "/home/samim/Documents/train.py", line 65, in <module>
pred = model.forward(i)
File "/home/samim/Documents/model.py", line 49, in forward
self.latent_rep = self.encoder(batch)
File "/home/samim/Documents/model.py", line 39, in encoder
self.kl_loss = kl_divergence(self.mu, self.log_sigma, dim=self.latent_dim)
File "/home/samim/Documents/model.py", line 59, in kl_divergence
out = 1 / 2. * torch.mean(torch.mean(1 / target_sigma**2 * means**2 + torch.exp(2 * log_sigma) / target_sigma**2 - 2 * log_sigma + 2 * torch.log(target_sigma), dim=1) - dim)
IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)
