I'm evaluating disentanglement in toy datasets seeing as we have such little understanding of the phenomena. I'm using various tools from differential geometry. Now I want to train a VAE on the following dataset:
class Radial(Dataset):
def __init__(self, num_samples=1000, num_radii=5, num_angles=8, r_noise=0.0, angle_noise=0.0):
"""
Initialize the dataset with random samples in polar coordinates.
The radial distances are equally spaced between 0.5 and 2.5.
The angles are equally spaced between 0 and 2 * pi.
Parameters:
- num_samples: Total number of samples in the dataset
- num_radii: Number of unique radial distances
- num_angles: Number of unique angles
- r_noise: Standard deviation of Gaussian noise added to the radial distances
- angle_noise: Standard deviation of Gaussian noise added to the angles
"""
self.num_samples = num_samples
self.num_radii = num_radii
self.num_angles = num_angles
# Generate the radial distances and angles
radii = np.linspace(0.5, 2.5, num_radii)
angles = np.linspace(0, 2 * np.pi, num_angles, endpoint=False)
# Create the grid of (r, theta)
r, theta = np.meshgrid(radii, angles)
# Flatten and repeat to create the dataset
r_flat = r.flatten()
theta_flat = theta.flatten()
repeats = num_samples // (num_radii * num_angles)
self.r = np.tile(r_flat, repeats)
self.theta = np.tile(theta_flat, repeats)
# Add Gaussian noise to the radial distances and angles
self.r += np.random.normal(0, r_noise, self.r.shape)
self.theta += np.random.normal(0, angle_noise, self.theta.shape)
# Convert to Cartesian coordinates (x, y)
self.x = self.r * np.cos(self.theta)
self.y = self.r * np.sin(self.theta)
# Convert to torch tensors
self.data = torch.tensor(np.column_stack((self.x, self.y)), dtype=torch.float32)
self.gen_factors = torch.tensor(np.column_stack((self.r, self.theta)), dtype=torch.float32)
def __len__(self):
return self.num_samples
def __getitem__(self, index):
return self.data[index]
The architecture that I came up with is as follows:
import torch
import torch.nn as nn
import torch.nn.functional as F
class Layer(nn.Module):
def __init__(self, input_dim, output_dim, act_func=nn.Tanh()):
super(Layer, self).__init__()
if act_func is None:
self.act_func = lambda x: x
else:
self.act_func = act_func
self.linear_map = nn.Linear(input_dim, output_dim)
self.out_features = output_dim
self.in_features = input_dim
def forward(self, x):
return self.act_func(self.linear_map(x))
class PolarLayer(nn.Module):
def __init__(self):
super(PolarLayer, self).__init__()
def forward(self, x):
r, theta = x.split(1, dim=-1)
x = r * torch.cos(theta)
y = r * torch.sin(theta)
return torch.cat([x, y], dim=-1)
class EuclLayer(nn.Module):
def __init__(self):
super(EuclLayer, self).__init__()
def forward(self, x):
# Assuming x is of shape [batch_size, 2]
# where x[:, 0] = radius (r) and x[:, 1] = angle (theta)
r = x[:, 0]
theta = x[:, 1]
# Convert to Cartesian coordinates
x_cartesian = r * torch.cos(theta)
y_cartesian = r * torch.sin(theta)
# Concatenate to form output tensor of shape [batch_size, 2]
output = torch.stack([x_cartesian, y_cartesian], dim=1)
return output
# Define the Encoder module
class Encoder(nn.Module):
def __init__(self, in_features, features, out_features):
super(Encoder, self).__init__()
self.in_features = in_features
self.out_features = out_features
self.layers = nn.ModuleList([
PolarLayer(), nn.Linear(in_features, features[0])
] + [
Layer(features[i], features[i + 1]) for i in range(len(features) - 1)
])
self.fc_mu = Layer(features[-1], out_features)
self.fc_log_var = Layer(features[-1], out_features)
def forward(self, x):
x = self.layers[0](x)
for layer in self.layers[1:]:
x = F.tanh(layer(x))
mu = self.fc_mu(x)
log_var = self.fc_log_var(x)
return mu, log_var
def forward_layers(self, x, indx):
assert indx < len(self.layers)
if indx == 0:
x = self.layers[0](x)
for layer in self.layers[indx:]:
x = layer.forward(x)
else:
for layer in self.layers[indx:]:
x = layer.forward(x)
return x
# Define the Decoder module
class Decoder(nn.Module):
def __init__(self, in_features, features, out_features):
super(Decoder, self).__init__()
self.in_features = in_features
self.out_features = out_features
self.layers = nn.ModuleList([
nn.Linear(in_features, features[0])
] + [
nn.Linear(features[i], features[i + 1]) for i in range(len(features) - 1)
] + [
nn.Linear(features[-1], out_features), EuclLayer()
])
def forward(self, x):
for layer in self.layers[:-1]:
x = F.tanh(layer(x))
x = self.layers[-1](x)
return x
class VAE(nn.Module):
def __init__(self, encoder, decoder):
super(VAE, self).__init__()
self.encoder = encoder
self.decoder = decoder
self.layers = encoder.layers
def reparameterize(self, mu, log_var):
std = torch.exp(0.5 * log_var)
eps = torch.randn_like(std)
return mu + eps * std
def forward(self, x):
mu, log_var = self.encoder(x)
z = self.reparameterize(mu, log_var)
x_reconstructed = self.decoder(z)
return x_reconstructed, mu, log_var
My main concern is whether or not to introduce polar coordinates. Without the conversion, my latent traversals for the radius look fine but on the angle it's simply a straight line, which is incorrect. Now I've tried a few different combinations for the architecture, tanh seems to work best for the activation function. However the reconstruction is always incorrect despite some very strong dampening of the KL divergence (for lambda = 0 to 0.5 through training, loss = recon_loss + lambda*KL). The traversals are sometimes better (I've included one that's on the poorer side) but they're not very consistent. Is there any other way of incorporating polar coordinates into the VAE?
The training script is below for reproducibility. Current default arguments lead to these plots.
from torch.optim.lr_scheduler import ReduceLROnPlateau
import matplotlib.pyplot as plt
import argparse
import os
from torch.utils.data import DataLoader
from torch.optim import Adam
from models.data.gen_factors import Radial
from models.unsupervised.vae.model import Encoder, Decoder, VAE
import torch.nn.functional as F
import torch
from tqdm import tqdm
def plot_reconstructions(model, dataset, save_dir, num_samples=10):
"""
Plot original and reconstructed data points.
"""
plt.figure(figsize=(12, 6))
# Sample data points from the dataset
sampled_data = [dataset[i] for i in range(num_samples)]
sampled_data = torch.stack(sampled_data)
# Get the reconstructed data points
with torch.no_grad():
reconstructed_data, _, _ = model(sampled_data)
plt.subplot(1, 2, 1)
plt.scatter(sampled_data[:, 0], sampled_data[:, 1], c='b', label='Original')
plt.title('Original Data Points')
plt.xlabel('x')
plt.ylabel('y')
plt.grid(True)
plt.legend()
plt.subplot(1, 2, 2)
plt.scatter(reconstructed_data[:, 0], reconstructed_data[:, 1], c='r', label='Reconstructed')
plt.title('Reconstructed Data Points')
plt.xlabel('x')
plt.ylabel('y')
plt.grid(True)
plt.legend()
plt.savefig(f"{save_dir}/recon.png")
def plot_latent_traversal(model, save_dir, dataset, num_points=10):
"""
Plot latent space traversal and observe how it affects the generated data.
"""
fig, ax = plt.subplots(1, 3, figsize=(18, 8))
# Retrieve the unique radii and angles from the dataset
unique_r = torch.unique(dataset.gen_factors[:, 0])
unique_theta = torch.unique(dataset.gen_factors[:, 1])
# Sample some unique radii and angles
sampled_r = unique_r[torch.randint(0, len(unique_r), (int(num_points**0.5), ))]
sampled_theta = unique_theta[torch.randint(0, len(unique_theta), (int(num_points**0.5), ))]
# Create a grid in the latent space corresponding to these radii and angles
z_r, z_theta = torch.meshgrid(sampled_r, sampled_theta)
# Flatten and create a tensor of shape [num_points * num_points, 2]
z = torch.stack([z_r.flatten(), z_theta.flatten()], dim=1)
# Decode the latent variables to generate data
with torch.no_grad():
generated_data = model.decoder(z)
# Plotting
ax[0].scatter(generated_data[:, 0], generated_data[:, 1], c='g', label='Generated Data')
ax[0].set_title('Latent Space Traversal')
ax[0].set_xlabel('x')
ax[0].set_ylabel('y')
ax[0].grid(True)
# Generate traversal values
traversal_values_radius = torch.linspace(0, 2.5, num_points)
travel_values_angles = torch.linspace(0, 2*torch.pi, num_points)
traversal_values = [traversal_values_radius, travel_values_angles]
for i in range(model.encoder.out_features):
z = torch.zeros(num_points, model.encoder.out_features)
z[:, i] = traversal_values[i]
with torch.no_grad():
generated_data = model.decoder(z)
ax[i+1].scatter(generated_data[:, 0], generated_data[:, 1], c='g', label=f'Latent Dim {i+1}')
ax[i+1].set_title(f'Latent Dimension {i+1} Traversal')
ax[i+1].set_xlabel('x')
ax[i+1].set_ylabel('y')
ax[i+1].grid(True)
plt.savefig(f"{save_dir}/latent_traversal.png")
# Define the VAE loss function (negative ELBO)
def vae_loss(x_reconstructed, x, mu, log_var):
# Reconstruction loss (MSE)
recon_loss = F.mse_loss(x_reconstructed, x, reduction='sum')
# KL divergence loss
kl_loss = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())
return recon_loss, kl_loss
def train(model, loader, criterion, optimizer, lambda_):
torch.manual_seed(args.seed)
model.train()
epoch_loss = 0
store_loss = []
store_recon_loss = []
store_kl_loss = []
for x in loader:
# Forward pass
x_reconstructed, mu, log_var = model(x)
# Compute loss
recon_loss, kl_loss = criterion(x_reconstructed, x, mu, log_var)
loss = recon_loss + lambda_*kl_loss
# Backward pass and optimization
optimizer.zero_grad()
loss.backward()
optimizer.step()
epoch_loss += loss.item()
store_loss.append(loss.item())
store_recon_loss.append(recon_loss.item())
store_kl_loss.append(lambda_*kl_loss.item())
return epoch_loss / len(loader.dataset), store_loss, store_recon_loss, store_kl_loss
def main(args):
model_name = f"R{args.num_radii}-A{args.num_angles}-Nr{args.radial_noise}-Na{args.angle_noise}"
save_dir = f"models/unsupervised/vae/saved_models/{model_name}"
plot_dir = f"models/unsupervised/vae/figures/{model_name}"
os.makedirs(save_dir, exist_ok=True)
os.makedirs(plot_dir, exist_ok=True)
# Initialize the VAE model
features = [128, 64]
encoder = Encoder(in_features=args.in_dim, features=features, out_features=args.out_dim)
decoder = Decoder(in_features=args.out_dim, features=list(reversed(features)), out_features=args.in_dim)
vae = VAE(encoder, decoder)
assert args.n_samples > args.num_radii*args.num_angles, "Number of samples must be greater than the number of radii times angles"
num_samples = args.n_samples - (args.n_samples % (args.num_radii*args.num_angles))
# Create the dataset and dataloader
if args.dataset == "radial":
dataset = Radial(num_samples, num_radii=args.num_radii, num_angles=args.num_angles, r_noise=args.radial_noise, angle_noise=args.angle_noise)
dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True)
# Initialize the optimizer
optimizer = Adam(vae.parameters(), lr=args.learning_rate)
# Training loop
num_epochs = args.num_epochs
store_loss = []
store_recon_loss = []
store_kl_loss = []
for epoch in range(num_epochs):
lambda_ = min(0.5, epoch/args.num_epochs)
epoch_loss, store_loss_, store_recon_loss_, store_kl_loss_ = train(vae, dataloader, vae_loss, optimizer, lambda_=lambda_)
store_loss += store_loss_
store_recon_loss += store_recon_loss_
store_kl_loss += store_kl_loss_
if epoch % args.SAVE_LOG == 0:
torch.save(vae.state_dict(), os.path.join(save_dir, f"model_{epoch}.pth"))
if args.verbose:
print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {epoch_loss:.4f}")
print("Training complete.")
if args.PLOT_LOSS:
plt.plot(store_loss)
plt.plot(store_recon_loss)
plt.plot(store_kl_loss)
plt.legend(["Total loss", "Reconstruction loss", "KL loss"])
plt.savefig(f"{plot_dir}/loss.png")
if args.PLOT_MODEL:
plot_reconstructions(vae, dataset, plot_dir, num_samples)
plot_latent_traversal(vae, plot_dir, dataset, num_samples)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--num_epochs', type=int, default=400)
parser.add_argument('--batch_size', type=int, default=128)
parser.add_argument('--learning_rate', type=float, default=0.01)
parser.add_argument('--out_dim', type=int, default=2)
parser.add_argument('--in_dim', type=int, default=2)
parser.add_argument('--n_samples', type=int, default=3000)
parser.add_argument('--dataset', type=str, default='radial')
parser.add_argument('--num_radii', type=int, default=5)
parser.add_argument('--num_angles', type=int, default=16)
parser.add_argument('--radial_noise', type=float, default=0.01)
parser.add_argument('--angle_noise', type=float, default=0.01)
parser.add_argument('--seed', type=int, default=2)
parser.add_argument('--PLOT_LOSS', type=bool, default=True)
parser.add_argument('--PLOT_MODEL', type=bool, default=True)
parser.add_argument('--SAVE_LOG', type=int, default=1)
parser.add_argument('--verbose', type=bool, default=True)
args = parser.parse_args()
main(args)
```
