def freeze_A_to_solution_and_fit():
# used to debug infs
# from tests.test_vi import test_elbo_components, test_q_E_logstick
SCALE = 1.
N = 500
X = generate_gg_blocks_dataset(N, 0.05)
model = InfiniteIBP(1.5, 6, 0.1, 0.05, 36)
model.phi.data[:4] = SCALE * gg_blocks()
model.init_z(N)
model.train()
visualize_A_save(model.phi.detach().numpy(), 0)
visualize_nu_save(model.nu.detach().numpy(), 0)
optimizer = torch.optim.Adam(model.parameters(), 0.1)
for i in range(20):
model.cavi(X)
print("[Epoch {:<3}] ELBO = {:.3f}".format(i + 1,
model.elbo(X).item()))
print("CHANGE OF REGIME")
visualize_A_save(model.phi.detach().numpy(), 20)
visualize_nu_save(model.nu.detach().numpy(), 20)
import ipdb
ipdb.set_trace()
def test_e_log_stick():
"""
This test DOES NOT PASS, and maybe should not
"""
model = InfiniteIBP(4., 10, 0.1, 0.5, 36)
model.init_z(10)
K = model.K
# take a lot of samples to get something working
dist = Beta(model.tau.detach()[:, 0], model.tau.detach()[:, 1])
samples = dist.sample((100000, ))
f = (1. - samples.cumprod(1)).log().mean(0)
log_stick, q = model._E_log_stick(model.tau, model.K)
jeffrey_q = np.zeros((K, K))
jeffrey_log_stick = np.zeros((K, ))
for k in range(K):
a, b = compute_q_Elogstick(model.tau.detach().numpy().T, k)
jeffrey_q[k, :k + 1] = a
jeffrey_log_stick[k] = b
print("old: {}".format(jeffrey_log_stick))
print("new: {}".format(log_stick.detach().numpy()))
print("samples: {}".format(f.detach().numpy()))
import ipdb
ipdb.set_trace()
def show_that_ADVI_init_doesnt_matter():
SCALE = 1.
N = 500
X = generate_gg_blocks_dataset(N, 0.05)
model = InfiniteIBP(1.5, 6, 0.1, 0.05, 36)
model.phi.data[:4] = SCALE * gg_blocks()
model.init_z(N)
model.train()
visualize_A_save(model.phi.detach().numpy(), 0)
visualize_nu_save(model.nu.detach().numpy(), 0)
optimizer = torch.optim.Adam(model.parameters(), 0.003)
for i in range(1000):
optimizer.zero_grad()
loss = -model.elbo(X)
print("[Epoch {:<3}] ELBO = {:.3f}".format(i + 1, -loss.item()))
loss.backward()
optimizer.step()
assert loss.item() != np.inf, "loss is inf"
visualize_A_save(model.phi.detach().numpy(), 1000)
visualize_nu_save(model.nu.detach().numpy(), 1000)
def freeze_A_to_solution_and_fit():
# used to debug infs
# from tests.test_vi import test_elbo_components, test_q_E_logstick
SCALE = 1.
N = 500
X = generate_gg_blocks_dataset(N, 0.05)
model = InfiniteIBP(1.5, 6, 0.1, 0.05, 36)
model.phi.data[:4] = SCALE * gg_blocks()
model.init_z(N)
model.train()
visualize_A_save(model.phi.detach().numpy(), 0)
visualize_nu_save(model.nu.detach().numpy(), 0)
optimizer = torch.optim.Adam(model.parameters(), 0.1)
for i in range(1000):
optimizer.zero_grad()
loss = -model.elbo(X)
print("[Epoch {:<3}] ELBO = {:.3f}".format(i + 1, -loss.item()))
loss.backward()
# zero out the grad on phi, phi_var
model.phi.grad.zero_()
model._phi_var.grad.zero_()
optimizer.step()
assert loss.item() != np.inf, "loss is inf"
print("CHANGE OF REGIME")
visualize_A_save(model.phi.detach().numpy(), 1000)
visualize_nu_save(model.nu.detach().numpy(), 1000)
for param_group in optimizer.param_groups:
param_group['lr'] = 0.001
for i in range(500):
optimizer.zero_grad()
loss = -model.elbo(X)
print("[Epoch {:<3}] ELBO = {:.3f}".format(i + 1, -loss.item()))
loss.backward()
optimizer.step()
assert loss.item() != np.inf, "loss is inf"
visualize_A_save(model.phi.detach().numpy(), 1500)
visualize_nu_save(model.nu.detach().numpy(), 1500)
import ipdb; ipdb.set_trace()
def check_that_naive_doesnt_work():
N = 500
X = generate_gg_blocks_dataset(N, 0.05)
model = InfiniteIBP(1.5, 6, 0.1, 0.05, 36)
model.init_z(N)
model.eval()
for i in range(100):
if i % 5 == 0:
visualize_A_save(model.phi.detach().numpy(), i)
if (i + 1) % 10 == 0:
model._nu.data = torch.randn(model._nu.shape)
model.cavi(X)
print("[Epoch {:<3}] ELBO = {:.3f}".format(i + 1,
model.elbo(X).item()))
def test_vectorized_cavi():
"""
OUT OF DATE (unnecessary?)
"""
model = InfiniteIBP(4., 6, 0.1, 0.05, 36)
model.init_z(10)
X = torch.randn(10, 36)
N, K, D = X.shape[0], model.K, model.D
slow_phi_var = torch.zeros(K, D)
slow_phi = torch.zeros(K, D)
for k in range(K):
precision = (1. / (model.sigma_a**2) + model.nu[:, k].sum() /
(model.sigma_n**2))
slow_phi_var[k] = torch.ones(model.D) / precision
s = 0
# this can definitely be parallelized - it's probably really slow right now
for n in range(N):
s += model.nu[n][k] * (X[n] -
(model.nu[n] @ model.phi - model.phi[k]))
slow_phi[k] = s / ((model.sigma_n**2) * precision)
fast_phi_var = torch.zeros(K, D)
fast_phi = torch.zeros(K, D)
for k in range(K):
precision = (1. / (model.sigma_a**2) + model.nu[:, k].sum() /
(model.sigma_n**2))
fast_phi_var[k] = torch.ones(model.D) / precision
s = (model.nu[:, k].view(N, 1) *
(X -
(model.nu @ model.phi - torch.ger(model.nu[:, k], model.phi[k])))
).sum(0)
fast_phi[k] = s / ((model.sigma_n**2) * precision)
print(slow_phi.detach().numpy())
print(fast_phi.detach().numpy())
assert (slow_phi - fast_phi).abs().max() < 1e-5, "Fast phi is wrong"
def find_a_better_scheme():
SCALE = 1.
N = 500
X = generate_gg_blocks_dataset(N, 0.05)
model = InfiniteIBP(1.5, 6, 0.1, 0.05, 36)
model.init_z(N)
model.train()
visualize_A_save(model.phi.detach().numpy(), 0)
visualize_nu_save(model.nu.detach().numpy(), 0)
optimizer = torch.optim.Adam([{
'params': [model._nu, model._tau]
}, {
'params': [model._phi_var, model.phi],
'lr': 0.003
}],
lr=0.1)
elbo_array = []
iter_count = 0
for j in range(6):
for i in range(1000):
optimizer.zero_grad()
loss = -model.elbo(X)
print("[Epoch {:<3}] ELBO = {:.3f}".format(i + 1, -loss.item()))
loss.backward()
optimizer.step()
iter_count += 1
assert loss.item() != np.inf, "loss is inf"
elbo_array.append(-loss.item())
visualize_A_save(model.phi.detach().numpy(), iter_count)
visualize_nu_save(model.nu.detach().numpy(), iter_count)
model._nu.data = torch.randn(model._nu.shape)
plt.plot(np.arange(len(elbo_array)), np.array(elbo_array))
plt.show()
import ipdb
ipdb.set_trace()
def grad_of_E_log_stick_wrt_q(inputs=None):
"""
So the value of q is set based on maximizing E_log_stick
"""
np.random.seed(0)
torch.manual_seed(0)
if inputs is None:
model = InfiniteIBP(4., 6, 0.1, 0.5, 36)
model.init_z(10)
model.train()
X = torch.randn(10, 36)
N, D = 10, 36
else:
model, X = inputs
N, D = X.shape[0]
_, q = model._E_log_stick(model.tau, model.K)
q.requires_grad = True
_E_log_stick = model._E_log_stick_from_q(q, model.tau)
_E_log_stick[1].backward()
import ipdb
ipdb.set_trace()
def check_gradient_of_tau_after_some_updates():
from matplotlib import pyplot as plt
SCALE = 1.
N = 500
X = generate_gg_blocks_dataset(N, 0.05)
model = InfiniteIBP(1.5, 6, 0.1, 0.05, 36)
model.init_z(N)
model.train()
optimizer = torch.optim.Adam([{'params': [model._nu, model._tau]},
{'params': [model._phi_var, model.phi], 'lr': 0.003}], lr=0.1)
elbo_array = []
iter_count = 0
# values = np.zeros((6 * 3000, 12))
for j in range(6):
for i in range(1000):
optimizer.zero_grad()
loss = -model.elbo(X)
print("[Epoch {:<3}] ELBO = {:.3f}".format(i + 1, -loss.item()))
loss.backward()
optimizer.step()
assert loss.item() != np.inf, "loss is inf"
elbo_array.append(-loss.item())
# values[iter_count] = model.tau.detach().numpy().reshape((-1,))
iter_count += 1
# plt.figure()
# for i in range(12):
# plt.plot(np.arange(3000), values[j*3000:(j + 1)*3000, i])
# plt.savefig('tau_set_{}.png'.format(j))
visualize_A_save(model.phi.detach().numpy(), iter_count)
visualize_nu_save(model.nu.detach().numpy(), iter_count)
model._nu.data = torch.randn(model._nu.shape)
def find_a_better_scheme(nu_resets=False,tempering=False):
from matplotlib import pyplot as plt
print("")
print("")
print("")
print("NU RESETS:",nu_resets)
print("TEMPERING:",tempering)
N = 500
sigma_n = 0.05
print("DATA SIZE:",N)
print("DATA NOISE:",sigma_n)
X = generate_gg_blocks_dataset(N, sigma_n)
alpha=1.5
K = 6
D = 36
sigma_a = 0.1
print("ALPHA:",alpha)
print("K:",K)
print("SIGMA N:",sigma_n)
print("SIGMA A:",sigma_a)
model = InfiniteIBP(alpha, K, sigma_a, sigma_n, D)
model.init_z(N)
if tempering:
print("INIT TEMPERING PARAMS")
M = 10
print("NUM TEMPERATURES:",M)
model.init_r_and_T(N,M)
model.train()
visualize_A_save(model.phi.detach().numpy(), 0)
visualize_nu_save(model.nu.detach().numpy(), 0)
if tempering:
print("Initing optimizer with tempering params included")
optimizer = torch.optim.Adam([{'params': [model._nu, model._tau, model._r]},
{'params': [model._phi_var, model.phi], 'lr': 0.003}], lr=0.1)
else:
optimizer = torch.optim.Adam([{'params': [model._nu, model._tau]},
{'params': [model._phi_var, model.phi], 'lr': 0.003}], lr=0.1)
elbo_array = []
iter_count = 0
for j in range(6):
for i in range(1000):
optimizer.zero_grad()
elbo = model.elbo(X)
loss = -elbo
if tempering:
loss = -model.elbo_tempered(X)
print("[Epoch {:<3}] ELBO = {:.3f}".format(i + 1, elbo.item()))
loss.backward()
optimizer.step()
iter_count += 1
assert loss.item() != np.inf, "loss is inf"
elbo_array.append(elbo.item())
visualize_A_save(model.phi.detach().numpy(), iter_count)
visualize_nu_save(model.nu.detach().numpy(), iter_count)
if nu_resets:
model._nu.data = torch.randn(model._nu.shape)
plt.plot(np.arange(len(elbo_array)), np.array(elbo_array))
plt.show()
import ipdb; ipdb.set_trace()
def infer_music():
# pip install librosa
# Load the audio data
# Compute the Spectrogram. N is time points. D is frequencies.
# It is DxN by default, so we transpose.
# We will do factorization for the magnitudes np.abs(spectrogram)
# We will need the phase np.angle(spectrogram) to resynthesize to audio
# Only keep first 500 time windows for now. To increase data size, keep all timesteps
y,sr = librosa.load('wav/bach.wav')
true_spec = librosa.stft(y,n_fft=2048)
true_mag = np.abs(true_spec)
true_mag = true_mag.T
true_mag = true_mag[:500,:]
true_phase = np.angle(true_spec)
true_phase = true_phase.T[:500,:].T
X=torch.from_numpy(true_mag)
N,D = X.size()
print("{} timepoints, {} frequencies".format(N,D))
# Don't know how many latents to do.
# Should be one for each frequency we hope to find,
# but can be less to find groups of frequencies.
# Using 30 for now
model = InfiniteIBP(1.5, 30, 0.1, 0.05, D)
model.init_z(N)
model.train()
optimizer = torch.optim.Adam([{'params': [model._nu, model._tau]},
{'params': [model._phi_var, model.phi], 'lr': 0.003}], lr=0.1)
elbo_array = []
iter_count = 0
# Should increase j range and i range
# Should increase the number of times nu is reset,
# but make sure not to reset it to close to when you stop doing inference
for j in range(10):
for i in range(1000):
optimizer.zero_grad()
loss = -model.elbo(X)
print("[Epoch {:<3}] ELBO = {:.3f}".format(i + 1, -loss.item()))
loss.backward()
optimizer.step()
iter_count += 1
assert loss.item() != np.inf, "loss is inf"
elbo_array.append(-loss.item())
#visualize_A_save(model.phi.detach().numpy(), iter_count)
#visualize_nu_save(model.nu.detach().numpy(), iter_count)
if j < 8:
model._nu.data = torch.randn(model._nu.shape)
plt.plot(np.arange(len(elbo_array)), np.array(elbo_array))
plt.show()
# Sample some assignments
# Take the mean of the features
# Recover the data, clip numbers below 0 because magnitude can't be negative
# Recover real and imaginary part of spectrogram from recovered magnitude and true phase
# Recover spectrogram from recovered real and imaginary parts.
# Resynthesize audio with istft
# When done, listen to bach_reconstruct.wav
p_Z = Bern(model.nu.clone())
Z = p_Z.sample()
A = model.phi.clone().detach()
mag = (Z@A).numpy().clip(min=0.0).T
real = np.multiply(mag,np.cos(true_phase))
imag = np.multiply(mag,np.sin(true_phase))
spec = real + imag*np.complex(0,1)
y_reconstruct = librosa.istft(spec)
librosa.output.write_wav('wav/bach_reconstruct.wav',y_reconstruct,sr)
import ipdb; ipdb.set_trace()
def test_cavi_updates_are_correct(inputs=None):
"""
PASSES for nu, phi
DOES NOT PASS for cavi_tau
Logic:
After doing the CAVI update for a single variable, we would expect that dELBO/dv is about 0.
"""
np.random.seed(0)
torch.manual_seed(0)
if inputs is None:
model = InfiniteIBP(4., 6, 0.1, 0.5, 36)
model.init_z(10)
model.train()
X = torch.randn(10, 36)
N, D = 10, 36
else:
model, X = inputs
N, D = X.shape[0]
optimizer = torch.optim.SGD(model.parameters(), 0.01)
optimizer.zero_grad()
# CAVI update for phi, _phi_var
k = 2
# model.eval()
# model.cavi_phi(k, X)
# model.train()
# optimizer.zero_grad()
# loss = model.elbo(X)
# loss.backward()
# assert model.phi.grad[k].abs().max().item() < 1e-4, "CAVI update for phi is wrong"
# assert model._phi_var.grad[k].abs().max().item() < 1e-4, "CAVI update for phi_var is wrong"
# optimizer.zero_grad()
# CAVI update for nu
# k = 3
# n = 2
# model.eval()
# log_stick, _ = model._E_log_stick(model.tau, model.K)
# model.cavi_nu(n, k, X, log_stick)
# # # compute the ELBO
# model.train()
# optimizer.zero_grad()
# loss = model.elbo(X)
# loss.backward()
# assert model._nu.grad[n][k].abs().max().item() < 1e-4, "CAVI update for nu is wrong"
# optimizer.zero_grad()
# CAVI update for tau
k = 0
model.eval()
_, q = model._E_log_stick(model.tau, model.K)
model.slow_cavi_tau(k, X, q)
# # compute the ELBO
model.train()
optimizer.zero_grad()
loss = model.elbo(X)
loss.backward()
print(model._tau.grad)
def test_q_E_logstick(inputs=None):
"""
Test the computation of the multinomial bound in Finale's tech report
by computing it by hand for k=2 and k=3
PASSES
"""
from scipy.special import digamma as dgm
if inputs is None:
K = 6
tau = 0.01 + nn.Softplus()(torch.randn(K, 2))
else:
tau, K = inputs
# # compute, by hand, the distribution q_2
hand_q = np.zeros(2)
taud = tau.clone().detach().numpy()
hand_q[0] = dgm(taud[0, 1]) - dgm(taud[0, 1] + taud[0, 0])
hand_q[1] = dgm(taud[1, 1]) + dgm(taud[0, 0]) - \
(dgm(taud[0, 0] + taud[0, 1]) + dgm(taud[1, 0] + taud[1, 1]))
hand_q = np.exp(hand_q)
hand_q /= hand_q.sum()
# let's check E_log_stick for k = 2 (i.e. index 1)
hand_E_logstick = 0
a = (hand_q[0] * dgm(taud[0, 1]) + hand_q[1] * dgm(taud[1, 1]))
b = (hand_q[1] * dgm(taud[0, 0]))
ca = dgm(taud[0, 0] + taud[0, 1]) * (hand_q[0] + hand_q[1])
cb = dgm(taud[1, 0] + taud[1, 1]) * (hand_q[1])
d = (hand_q[0] * np.log(hand_q[0]) + hand_q[1] * np.log(hand_q[1]))
hand_E_logstick = a + b - ca - cb - d
# print(a, b, ca + cb, d)
# for q_3
hand_q_3 = np.zeros(3)
hand_q_3[0] = dgm(taud[0, 1]) - dgm(taud[0, 1] + taud[0, 0])
hand_q_3[1] = dgm(taud[1, 1]) + dgm(taud[0, 0]) - \
(dgm(taud[0, 0] + taud[0, 1]) + dgm(taud[1, 0] + taud[1, 1]))
hand_q_3[2] = dgm(taud[2, 1]) + (dgm(taud[0, 0]) + dgm(taud[1, 0])) - \
(dgm(taud[0, 0] + taud[0, 1]) + dgm(taud[1, 0] + taud[1, 1]) + dgm(taud[2, 0] + taud[2, 1]))
hand_q_3 = np.exp(hand_q_3)
hand_q_3 /= hand_q_3.sum()
# let's check E_log_stick for k = 3 (i.e. index 2)
hand_E_logstick_3 = 0
hand_E_logstick_3 += (hand_q_3[0] * dgm(taud[0, 1]) +
hand_q_3[1] * dgm(taud[1, 1]) +
hand_q_3[2] * dgm(taud[2, 1]))
hand_E_logstick_3 += ((hand_q_3[1] + hand_q_3[2]) * dgm(taud[0, 0])) + \
((hand_q_3[2]) * dgm(taud[1, 0]))
hand_E_logstick_3 -= (dgm(taud[0, 0] + taud[0, 1]) * (hand_q_3[0] + hand_q_3[1] + hand_q_3[2]) + \
dgm(taud[1, 0] + taud[1, 1]) * (hand_q_3[1] + hand_q_3[2]) + \
dgm(taud[2, 0] + taud[2, 1]) * (hand_q_3[2]))
hand_E_logstick_3 -= (hand_q_3[0] * np.log(hand_q_3[0]) + \
hand_q_3[1] * np.log(hand_q_3[1]) + \
hand_q_3[2] * np.log(hand_q_3[2]))
# since the errors seem to appear later down inthe term, let's check the value of _E_log_stick[4] but use the (hopefully correct) q
E_logstick, q = InfiniteIBP._E_log_stick(tau, K)
k = 3 # this is the index, so k = 4 in the equations (i.e. 1, 2, 3, 4), so q is a 4-dim Cat
hand_E_logstick_4 = 0
for m in range(k + 1):
hand_E_logstick_4 += q[k][m] * digamma(tau[m, 1])
for m in range(k):
q_sum = 0
for n in range(m + 1, k + 1):
q_sum += q[k][n]
hand_E_logstick_4 += q_sum * digamma(tau[m, 0])
for m in range(k + 1):
q_sum = 0
for n in range(m, k + 1):
q_sum += q[k][n]
hand_E_logstick_4 -= q_sum * digamma(tau[m, 0] + tau[m, 1])
for m in range(k + 1):
hand_E_logstick_4 -= q[k][m] * (q[k][m] + EPS).log()
# test that the computed q is equal to the hand-computed q
assert np.abs(
(q[1, :2].numpy() -
hand_q)).max() < 1e-6, "_E_log_stick doesn't compute q_2 correctly"
assert np.abs(
(q[2, :3].numpy() -
hand_q_3)).max() < 1e-6, "_E_log_stick doesn't compute q_3 correctly"
# test that the hand-computed e logstick is equal to the computed e logstick
assert np.abs(E_logstick[1].item() - hand_E_logstick).max(
) < 1e-5, "_E_logstick_2 isn't computed correctly"
assert np.abs(E_logstick[2].item() - hand_E_logstick_3).max(
) < 1e-5, "_E_logstick_3 isn't computed correctly"
assert np.abs(E_logstick[3].item() - hand_E_logstick_4.item()).max(
) < 1e-5, "_E_logstick_4 isn't computed correctly"
def test_elbo_components(inputs=None):
"""
Test that various KL divergences are positive, and in the case of the
approximate posterior q(v), compute it exactly in two ways and check
that both give the same result.
"""
if inputs is None:
model = InfiniteIBP(4., 6, 0.1, 0.5, 36)
model.init_z(10)
model.train()
X = torch.randn(10, 36)
else:
model, X = inputs
a = model._1_feature_prob(model.tau).sum()
b = model._2_feature_assign(model.nu, model.tau).sum()
c = model._3_feature_prob(model.phi_var, model.phi).sum()
d = model._4_likelihood(X, model.nu, model.phi_var, model.phi).sum()
e = model._5_entropy(model.tau, model.phi_var, model.nu).sum()
entropy_q_v = InfiniteIBP._entropy_q_v(model.tau)
entropy_q_A = InfiniteIBP._entropy_q_A(model.phi_var)
entropy_q_z = InfiniteIBP._entropy_q_z(model.nu)
try:
assert (a + b + c + d + e).item() not in (np.inf,
-np.inf), "ELBO is inf"
except AssertionError:
print("a: ", a)
print("b: ", b)
print("c: ", c)
print("d: ", d)
print("e: ", e)
print("entropy_q_v: ", entropy_q_v)
print("entropy_q_A: ", entropy_q_A)
print("entropy_q_z: ", entropy_q_z)
raise
# check the sign of the various KL divergences (summed, so less powerful than it could be)
assert (a + entropy_q_v).item() <= 0, "KL(q(pi) || p(pi)) is negative"
# assert (b + entropy_q_z).item() <= 10, "KL(q(z) || p(z)) is negative" # we give this one some tolerance
assert (c + entropy_q_A).item() <= 0, "KL(q(A) || p(A)) is negative"
assert (a + b + c + e
).item() <= 0, "KL divergence between q(...) || p(...) is negative"
# check the empirical value of the component KL divergences (this is a very strong test)
from torch.distributions import Beta, kl_divergence
p_pi = Beta(model.alpha, 1.)
q_pi = Beta(model.tau[:, 0], model.tau[:, 1])
try:
assert (
kl_divergence(q_pi, p_pi).sum() +
(a + entropy_q_v)).abs() < 1e-3, "KL(q(pi) || p(pi)) is incorrect"
except:
import ipdb
ipdb.set_trace()
import sys, os
sys.path.insert(0, os.path.abspath('..'))
# inside-package imports below here
from src.vi import InfiniteIBP
from src.utils import register_hooks, visualize_A, visualize_A_save, visualize_nu_save
from src.data import generate_gg_blocks, generate_gg_blocks_dataset, gg_blocks
N = 500
X = generate_gg_blocks_dataset(N, 0.05)
model = InfiniteIBP(1.5, 6, 0.1, 0.05, 36)
model.init_z(N)
model.train()
visualize_A_save(model.phi.detach().numpy(), 0)
visualize_nu_save(model.nu.detach().numpy(), 0)
optimizer = torch.optim.Adam([{'params': [model._nu, model._tau]},
{'params': [model._phi_var, model.phi], 'lr': 0.003}], lr=0.1)
elbo_array = []
iter_count = 0
for j in range(10):
for i in range(1000):
optimizer.zero_grad()
loss = -model.elbo(X)