Conditional Normalizing Flow Model¶
Here, we train a conditional normalizing flow model $q(x|c)$. Our target $p(x|c)$ is a simple 2D Gaussian $\mathcal{N}(x|\mu, \sigma)$, where we condition on the mean $\mu$ and standard deviation $\sigma$, i.e. $c = (\mu, \sigma)$. We apply conditional autoregressive and coupling neural spline flows as well as a conditional masked autoregressive flow to the problem.
Setup¶
In [1]:
Copied!
# Import packages
import torch
import numpy as np
import normflows as nf
from matplotlib import pyplot as plt
from tqdm import tqdm
# Import packages
import torch
import numpy as np
import normflows as nf
from matplotlib import pyplot as plt
from tqdm import tqdm
In [2]:
Copied!
# Get device to be used
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# Get device to be used
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
In [3]:
Copied!
# Define target
target = nf.distributions.target.ConditionalDiagGaussian()
context_size = 4
# Plot target
grid_size = 100
xx, yy = torch.meshgrid(torch.linspace(-2, 2, grid_size), torch.linspace(-2, 2, grid_size), indexing='ij')
zz = torch.cat([xx.unsqueeze(2), yy.unsqueeze(2)], 2).view(-1, 2)
zz = zz.to(device)
context_plot = torch.cat([torch.tensor([0.3, 0.9]).to(device) + torch.zeros_like(zz),
0.6 * torch.ones_like(zz)], dim=-1)
logp = target.log_prob(zz, context_plot)
p_target = torch.exp(logp).view(*xx.shape).cpu().data.numpy()
plt.figure(figsize=(10, 10))
plt.pcolormesh(xx, yy, p_target, shading='auto')
plt.gca().set_aspect('equal', 'box')
plt.show()
# Define target
target = nf.distributions.target.ConditionalDiagGaussian()
context_size = 4
# Plot target
grid_size = 100
xx, yy = torch.meshgrid(torch.linspace(-2, 2, grid_size), torch.linspace(-2, 2, grid_size), indexing='ij')
zz = torch.cat([xx.unsqueeze(2), yy.unsqueeze(2)], 2).view(-1, 2)
zz = zz.to(device)
context_plot = torch.cat([torch.tensor([0.3, 0.9]).to(device) + torch.zeros_like(zz),
0.6 * torch.ones_like(zz)], dim=-1)
logp = target.log_prob(zz, context_plot)
p_target = torch.exp(logp).view(*xx.shape).cpu().data.numpy()
plt.figure(figsize=(10, 10))
plt.pcolormesh(xx, yy, p_target, shading='auto')
plt.gca().set_aspect('equal', 'box')
plt.show()
Autoregressive Neural Spline Flow¶
Model specification¶
In [4]:
Copied!
# Define flows
K = 4
latent_size = 2
hidden_units = 128
hidden_layers = 2
flows = []
for i in range(K):
flows += [nf.flows.AutoregressiveRationalQuadraticSpline(latent_size, hidden_layers, hidden_units,
num_context_channels=context_size)]
flows += [nf.flows.LULinearPermute(latent_size)]
# Set base distribution
q0 = nf.distributions.DiagGaussian(2, trainable=False)
# Construct flow model
model = nf.ConditionalNormalizingFlow(q0, flows, target)
# Move model on GPU if available
model = model.to(device)
# Define flows
K = 4
latent_size = 2
hidden_units = 128
hidden_layers = 2
flows = []
for i in range(K):
flows += [nf.flows.AutoregressiveRationalQuadraticSpline(latent_size, hidden_layers, hidden_units,
num_context_channels=context_size)]
flows += [nf.flows.LULinearPermute(latent_size)]
# Set base distribution
q0 = nf.distributions.DiagGaussian(2, trainable=False)
# Construct flow model
model = nf.ConditionalNormalizingFlow(q0, flows, target)
# Move model on GPU if available
model = model.to(device)
In [5]:
Copied!
# Plot initial flow distribution, target as contours
model.eval()
log_prob = model.log_prob(zz, context_plot).to('cpu').view(*xx.shape)
model.train()
prob = torch.exp(log_prob)
prob[torch.isnan(prob)] = 0
plt.figure(figsize=(10, 10))
plt.pcolormesh(xx, yy, prob.data.numpy(), shading='auto')
plt.contour(xx, yy, p_target, cmap=plt.get_cmap('cool'), linewidths=2)
plt.gca().set_aspect('equal', 'box')
plt.show()
# Plot initial flow distribution, target as contours
model.eval()
log_prob = model.log_prob(zz, context_plot).to('cpu').view(*xx.shape)
model.train()
prob = torch.exp(log_prob)
prob[torch.isnan(prob)] = 0
plt.figure(figsize=(10, 10))
plt.pcolormesh(xx, yy, prob.data.numpy(), shading='auto')
plt.contour(xx, yy, p_target, cmap=plt.get_cmap('cool'), linewidths=2)
plt.gca().set_aspect('equal', 'box')
plt.show()
Training¶
In [6]:
Copied!
# Train model
max_iter = 5000
batch_size= 128
loss_hist = np.array([])
optimizer = torch.optim.Adam(model.parameters(), lr=3e-4, weight_decay=1e-5)
for it in tqdm(range(max_iter)):
optimizer.zero_grad()
# Get training samples
context = torch.cat([torch.randn((batch_size, 2), device=device),
0.5 + 0.5 * torch.rand((batch_size, 2), device=device)],
dim=-1)
x = target.sample(batch_size, context)
# Compute loss
loss = model.forward_kld(x, context)
# Do backprop and optimizer step
if ~(torch.isnan(loss) | torch.isinf(loss)):
loss.backward()
optimizer.step()
# Log loss
loss_hist = np.append(loss_hist, loss.to('cpu').data.numpy())
# Plot loss
plt.figure(figsize=(10, 10))
plt.plot(loss_hist, label='loss')
plt.legend()
plt.show()
# Train model
max_iter = 5000
batch_size= 128
loss_hist = np.array([])
optimizer = torch.optim.Adam(model.parameters(), lr=3e-4, weight_decay=1e-5)
for it in tqdm(range(max_iter)):
optimizer.zero_grad()
# Get training samples
context = torch.cat([torch.randn((batch_size, 2), device=device),
0.5 + 0.5 * torch.rand((batch_size, 2), device=device)],
dim=-1)
x = target.sample(batch_size, context)
# Compute loss
loss = model.forward_kld(x, context)
# Do backprop and optimizer step
if ~(torch.isnan(loss) | torch.isinf(loss)):
loss.backward()
optimizer.step()
# Log loss
loss_hist = np.append(loss_hist, loss.to('cpu').data.numpy())
# Plot loss
plt.figure(figsize=(10, 10))
plt.plot(loss_hist, label='loss')
plt.legend()
plt.show()
100%|████████████████████████████████████████████████████████████| 5000/5000 [01:34<00:00, 52.69it/s]
Evaluation¶
In [7]:
Copied!
# Plot trained flow distribution, target as contours
model.eval()
log_prob = model.log_prob(zz, context_plot).to('cpu').view(*xx.shape)
model.train()
prob = torch.exp(log_prob)
prob[torch.isnan(prob)] = 0
plt.figure(figsize=(10, 10))
plt.pcolormesh(xx, yy, prob.data.numpy(), shading='auto')
plt.contour(xx, yy, p_target, cmap=plt.get_cmap('cool'), linewidths=2)
plt.gca().set_aspect('equal', 'box')
plt.show()
# Plot trained flow distribution, target as contours
model.eval()
log_prob = model.log_prob(zz, context_plot).to('cpu').view(*xx.shape)
model.train()
prob = torch.exp(log_prob)
prob[torch.isnan(prob)] = 0
plt.figure(figsize=(10, 10))
plt.pcolormesh(xx, yy, prob.data.numpy(), shading='auto')
plt.contour(xx, yy, p_target, cmap=plt.get_cmap('cool'), linewidths=2)
plt.gca().set_aspect('equal', 'box')
plt.show()
Coupling Neural Spline Flow¶
Model specification¶
In [8]:
Copied!
# Define flows
K = 4
latent_size = 2
hidden_units = 128
hidden_layers = 2
flows = []
for i in range(K):
flows += [nf.flows.CoupledRationalQuadraticSpline(latent_size, hidden_layers, hidden_units,
num_context_channels=context_size)]
flows += [nf.flows.LULinearPermute(latent_size)]
# Set base distribution
q0 = nf.distributions.DiagGaussian(2, trainable=False)
# Construct flow model
model = nf.ConditionalNormalizingFlow(q0, flows, target)
# Move model on GPU if available
model = model.to(device)
# Define flows
K = 4
latent_size = 2
hidden_units = 128
hidden_layers = 2
flows = []
for i in range(K):
flows += [nf.flows.CoupledRationalQuadraticSpline(latent_size, hidden_layers, hidden_units,
num_context_channels=context_size)]
flows += [nf.flows.LULinearPermute(latent_size)]
# Set base distribution
q0 = nf.distributions.DiagGaussian(2, trainable=False)
# Construct flow model
model = nf.ConditionalNormalizingFlow(q0, flows, target)
# Move model on GPU if available
model = model.to(device)
In [9]:
Copied!
# Plot initial flow distribution, target as contours
model.eval()
log_prob = model.log_prob(zz, context_plot).to('cpu').view(*xx.shape)
model.train()
prob = torch.exp(log_prob)
prob[torch.isnan(prob)] = 0
plt.figure(figsize=(10, 10))
plt.pcolormesh(xx, yy, prob.data.numpy(), shading='auto')
plt.contour(xx, yy, p_target, cmap=plt.get_cmap('cool'), linewidths=2)
plt.gca().set_aspect('equal', 'box')
plt.show()
# Plot initial flow distribution, target as contours
model.eval()
log_prob = model.log_prob(zz, context_plot).to('cpu').view(*xx.shape)
model.train()
prob = torch.exp(log_prob)
prob[torch.isnan(prob)] = 0
plt.figure(figsize=(10, 10))
plt.pcolormesh(xx, yy, prob.data.numpy(), shading='auto')
plt.contour(xx, yy, p_target, cmap=plt.get_cmap('cool'), linewidths=2)
plt.gca().set_aspect('equal', 'box')
plt.show()
Training¶
In [10]:
Copied!
# Train model
max_iter = 5000
batch_size= 128
loss_hist = np.array([])
optimizer = torch.optim.Adam(model.parameters(), lr=3e-4, weight_decay=1e-5)
for it in tqdm(range(max_iter)):
optimizer.zero_grad()
# Get training samples
context = torch.cat([torch.randn((batch_size, 2), device=device),
0.5 + 0.5 * torch.rand((batch_size, 2), device=device)],
dim=-1)
x = target.sample(batch_size, context)
# Compute loss
loss = model.forward_kld(x, context)
# Do backprop and optimizer step
if ~(torch.isnan(loss) | torch.isinf(loss)):
loss.backward()
optimizer.step()
# Log loss
loss_hist = np.append(loss_hist, loss.to('cpu').data.numpy())
# Plot loss
plt.figure(figsize=(10, 10))
plt.plot(loss_hist, label='loss')
plt.legend()
plt.show()
# Train model
max_iter = 5000
batch_size= 128
loss_hist = np.array([])
optimizer = torch.optim.Adam(model.parameters(), lr=3e-4, weight_decay=1e-5)
for it in tqdm(range(max_iter)):
optimizer.zero_grad()
# Get training samples
context = torch.cat([torch.randn((batch_size, 2), device=device),
0.5 + 0.5 * torch.rand((batch_size, 2), device=device)],
dim=-1)
x = target.sample(batch_size, context)
# Compute loss
loss = model.forward_kld(x, context)
# Do backprop and optimizer step
if ~(torch.isnan(loss) | torch.isinf(loss)):
loss.backward()
optimizer.step()
# Log loss
loss_hist = np.append(loss_hist, loss.to('cpu').data.numpy())
# Plot loss
plt.figure(figsize=(10, 10))
plt.plot(loss_hist, label='loss')
plt.legend()
plt.show()
100%|████████████████████████████████████████████████████████████| 5000/5000 [02:16<00:00, 36.51it/s]
Evaluation¶
In [11]:
Copied!
# Plot trained flow distribution, target as contours
model.eval()
log_prob = model.log_prob(zz, context_plot).to('cpu').view(*xx.shape)
model.train()
prob = torch.exp(log_prob)
prob[torch.isnan(prob)] = 0
plt.figure(figsize=(10, 10))
plt.pcolormesh(xx, yy, prob.data.numpy(), shading='auto')
plt.contour(xx, yy, p_target, cmap=plt.get_cmap('cool'), linewidths=2)
plt.gca().set_aspect('equal', 'box')
plt.show()
# Plot trained flow distribution, target as contours
model.eval()
log_prob = model.log_prob(zz, context_plot).to('cpu').view(*xx.shape)
model.train()
prob = torch.exp(log_prob)
prob[torch.isnan(prob)] = 0
plt.figure(figsize=(10, 10))
plt.pcolormesh(xx, yy, prob.data.numpy(), shading='auto')
plt.contour(xx, yy, p_target, cmap=plt.get_cmap('cool'), linewidths=2)
plt.gca().set_aspect('equal', 'box')
plt.show()
Masked Autoregressive Flow¶
Model specification¶
In [12]:
Copied!
# Define flows
K = 4
latent_size = 2
hidden_units = 128
num_blocks = 2
flows = []
for i in range(K):
flows += [nf.flows.MaskedAffineAutoregressive(latent_size, hidden_units,
context_features=context_size,
num_blocks=num_blocks)]
flows += [nf.flows.LULinearPermute(latent_size)]
# Set base distribution
q0 = nf.distributions.DiagGaussian(2, trainable=False)
# Construct flow model
model = nf.ConditionalNormalizingFlow(q0, flows, target)
# Move model on GPU if available
model = model.to(device)
# Define flows
K = 4
latent_size = 2
hidden_units = 128
num_blocks = 2
flows = []
for i in range(K):
flows += [nf.flows.MaskedAffineAutoregressive(latent_size, hidden_units,
context_features=context_size,
num_blocks=num_blocks)]
flows += [nf.flows.LULinearPermute(latent_size)]
# Set base distribution
q0 = nf.distributions.DiagGaussian(2, trainable=False)
# Construct flow model
model = nf.ConditionalNormalizingFlow(q0, flows, target)
# Move model on GPU if available
model = model.to(device)
In [13]:
Copied!
# Plot initial flow distribution, target as contours
model.eval()
log_prob = model.log_prob(zz, context_plot).to('cpu').view(*xx.shape)
model.train()
prob = torch.exp(log_prob)
prob[torch.isnan(prob)] = 0
plt.figure(figsize=(10, 10))
plt.pcolormesh(xx, yy, prob.data.numpy(), shading='auto')
plt.contour(xx, yy, p_target, cmap=plt.get_cmap('cool'), linewidths=2)
plt.gca().set_aspect('equal', 'box')
plt.show()
# Plot initial flow distribution, target as contours
model.eval()
log_prob = model.log_prob(zz, context_plot).to('cpu').view(*xx.shape)
model.train()
prob = torch.exp(log_prob)
prob[torch.isnan(prob)] = 0
plt.figure(figsize=(10, 10))
plt.pcolormesh(xx, yy, prob.data.numpy(), shading='auto')
plt.contour(xx, yy, p_target, cmap=plt.get_cmap('cool'), linewidths=2)
plt.gca().set_aspect('equal', 'box')
plt.show()
Training¶
In [14]:
Copied!
# Train model
max_iter = 5000
batch_size= 128
loss_hist = np.array([])
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-5)
for it in tqdm(range(max_iter)):
optimizer.zero_grad()
# Get training samples
context = torch.cat([torch.randn((batch_size, 2), device=device),
0.5 + 0.5 * torch.rand((batch_size, 2), device=device)],
dim=-1)
x = target.sample(batch_size, context)
# Compute loss
loss = model.forward_kld(x, context)
# Do backprop and optimizer step
if ~(torch.isnan(loss) | torch.isinf(loss)):
loss.backward()
optimizer.step()
# Log loss
loss_hist = np.append(loss_hist, loss.to('cpu').data.numpy())
# Plot loss
plt.figure(figsize=(10, 10))
plt.plot(loss_hist, label='loss')
plt.legend()
plt.show()
# Train model
max_iter = 5000
batch_size= 128
loss_hist = np.array([])
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-5)
for it in tqdm(range(max_iter)):
optimizer.zero_grad()
# Get training samples
context = torch.cat([torch.randn((batch_size, 2), device=device),
0.5 + 0.5 * torch.rand((batch_size, 2), device=device)],
dim=-1)
x = target.sample(batch_size, context)
# Compute loss
loss = model.forward_kld(x, context)
# Do backprop and optimizer step
if ~(torch.isnan(loss) | torch.isinf(loss)):
loss.backward()
optimizer.step()
# Log loss
loss_hist = np.append(loss_hist, loss.to('cpu').data.numpy())
# Plot loss
plt.figure(figsize=(10, 10))
plt.plot(loss_hist, label='loss')
plt.legend()
plt.show()
100%|████████████████████████████████████████████████████████████| 5000/5000 [02:00<00:00, 41.53it/s]
Evaluation¶
In [15]:
Copied!
# Plot trained flow distribution, target as contours
model.eval()
log_prob = model.log_prob(zz, context_plot).to('cpu').view(*xx.shape)
model.train()
prob = torch.exp(log_prob)
prob[torch.isnan(prob)] = 0
plt.figure(figsize=(10, 10))
plt.pcolormesh(xx, yy, prob.data.numpy(), shading='auto')
plt.contour(xx, yy, p_target, cmap=plt.get_cmap('cool'), linewidths=2)
plt.gca().set_aspect('equal', 'box')
plt.show()
# Plot trained flow distribution, target as contours
model.eval()
log_prob = model.log_prob(zz, context_plot).to('cpu').view(*xx.shape)
model.train()
prob = torch.exp(log_prob)
prob[torch.isnan(prob)] = 0
plt.figure(figsize=(10, 10))
plt.pcolormesh(xx, yy, prob.data.numpy(), shading='auto')
plt.contour(xx, yy, p_target, cmap=plt.get_cmap('cool'), linewidths=2)
plt.gca().set_aspect('equal', 'box')
plt.show()
