How to train VAE to learn the distribution of 3 classes - pytorch

Question

I want to train a Variational Auto-Encoder for a simple task: learn 3-class dataset and generate new samples from it's latent space. The problem is that despite it seems like an easy task I can't make it work. I expect it to produce the same 3 clusters when generating new samples from decoder (2 image).
I played with different values for beta. It's doesn't make it better.
How I can fix it and get the desired output or there some restrictions? Thanks in advance.
It's the code below, you may easy copypast the 3 code pieces into 3 cells in notebook and run.
Imports:
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import minimize
import torch
from torch.utils.data import Dataset
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from sklearn.datasets import make_classification
Here's dataset:
class TrainDataset(Dataset):
def __init__(self, X):
self.data = X.astype('float32')
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
return self.data[idx], 0
X, y = make_classification(n_samples=100, n_features=3, n_informative=3, n_redundant=0, n_classes=3, n_clusters_per_class=1, class_sep=2, random_state=16)
train_ratio = 0.7
batch_size = 8
train_loader = torch.utils.data.DataLoader(TrainDataset(X[:int(train_ratio * len(X))]), batch_size=batch_size)
test_loader = torch.utils.data.DataLoader(TrainDataset(X[int(train_ratio * len(X)):]), batch_size=batch_size)
fig = plt.figure(figsize=(10, 10))
ax = fig.add_subplot(projection='3d')
ax.scatter(X[:, 0], X[:, 1], X[:, 2], c=y)
plt.show()
Now the VAE:
class VAEmini(nn.Module):
def __init__(self, latent_d=2):
super().__init__()
self.fc1 = nn.Linear(3, 3)
self.fc21 = nn.Linear(3, latent_d)
self.fc22 = nn.Linear(3, latent_d)
self.fc3 = nn.Linear(latent_d, 3)
self.fc4 = nn.Linear(3, 3)
def encode(self, x):
h1 = F.relu(self.fc1(x))
return self.fc21(h1), self.fc22(h1)
def reparameterize(self, mu, logvar):
std = torch.exp(0.5*logvar)
eps = torch.randn_like(std)
return mu + eps*std
def decode(self, z):
h3 = F.relu(self.fc3(z))
return self.fc4(h3)
def forward(self, x):
mu, logvar = self.encode(x)
z = self.reparameterize(mu, logvar)
return self.decode(z), mu, logvar
def loss_function(recon_x, x, mu, logvar, beta):
BCE = F.mse_loss(recon_x, x, reduction='mean')
KLD = -0.5 * torch.mean(1 + logvar - mu.pow(2) - logvar.exp())
return BCE + beta * KLD
def train(epoch, beta):
model_mini.train()
train_loss = 0
for batch_idx, (data, _) in enumerate(train_loader):
data = data.to(device)
optimizer.zero_grad()
recon_batch, mu, logvar = model_mini(data)
loss = loss_function(recon_batch, data, mu, logvar, beta)
loss.backward()
train_loss += loss.item()
optimizer.step()
if batch_idx % log_interval == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader),
loss.item() / len(data)))
print('====> Epoch: {} Average loss: {:.4f}'.format(
epoch, train_loss / len(train_loader.dataset)))
def test(epoch, beta):
model_mini.eval()
test_loss = 0
with torch.no_grad():
for i, (data, _) in enumerate(test_loader):
data = data.to(device)
recon_batch, mu, logvar = model_mini(data)
test_loss += loss_function(recon_batch, data, mu, logvar, beta).item()
print('====> Test set loss: {:.4f}'.format(test_loss / len(test_loader.dataset)))
return test_loss / len(test_loader.dataset)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_mini = VAEmini(latent_d=2).to(device)
log_interval = 5
epochs = 100
optimizer = optim.Adam(model_mini.parameters(), lr=1e-2)
beta=0.01
for epoch in range(1, epochs + 1):
train(epoch, beta=beta)
test(epoch, beta=beta)
with torch.no_grad():
sample = torch.randn(1000, 2).to(device)
sample = model_mini.decode(sample).cpu()
fig = plt.figure(figsize=(20, 20))
ax = fig.add_subplot(projection='3d')
ax.scatter(sample[:, 0], sample[:, 1], sample[:, 2])
plt.show()