def make_nf(n_blocks, base_dist, affine_flow=StructuredAffineFlow):
blocks = []
for _ in range(n_blocks):
blocks += [affine_flow(2), PReLUFlow(2), BatchNormFlow(2)]
blocks += [affine_flow(2)]
return NormalizingFlow(
*blocks,
base_dist=base_dist,
)
def make_nf(n_blocks, affine_class, base_dist):
blocks = []
for _ in range(n_blocks):
blocks += [affine_class(2), PReLUFlow(2), BatchNormFlow(2)]
blocks += [affine_class(2)]
return NormalizingFlow(
*blocks,
base_dist=base_dist,
)
def make_nf(shared_blocks, n_blocks, base_dist):
blocks = [] + shared_blocks
for _ in range(n_blocks):
blocks += [StructuredAffineFlow(2), PReLUFlow(2), BatchNormFlow(2)]
blocks += [StructuredAffineFlow(2)]
# blocks += [AffineLUFlow(2), PReLUFlow(2), BatchNormFlow(2)]
# blocks += [AffineLUFlow(2)]
return NormalizingFlow(
*blocks,
base_dist=base_dist,
)
idx_0 = np.logical_and(X0[:, 0] < 0, X0[:, 1] < 0)
colors[idx_0] = 0
idx_1 = np.logical_and(X0[:, 0] >= 0, X0[:, 1] < 0)
colors[idx_1] = 1
idx_2 = np.logical_and(X0[:, 0] >= 0, X0[:, 1] >= 0)
colors[idx_2] = 2
idx_3 = np.logical_and(X0[:, 0] < 0, X0[:, 1] >= 0)
colors[idx_3] = 3
# %%
plt.scatter(X0[:, 0], X0[:, 1], s=5, c=colors)
# %%
blocks = sum(
[[StructuredAffineFlow(2), PReLUFlow(2)] for _ in range(5)] + [[StructuredAffineFlow(2)]],
[])
# %%
flow = NormalizingFlow(
*blocks,
base_dist=base_dist,
)
opt = optim.Adam(flow.parameters(), lr=5e-4)
# %%
count_parameters(flow)
# %%
writer = SummaryWriter(f"/workspace/sandbox/tensorboard_logs/{now_str()}")
# blocks += [AffineLUFlow(2)]
return NormalizingFlow(
*blocks,
base_dist=base_dist,
)
# %%
xdim = 2
hdim = 3
n_hidden = 3
n_classes = 3
shared_blocks = sum(
[[StructuredAffineFlow(2),
PReLUFlow(2), BatchNormFlow(2)] for _ in range(3)], [])
#shared_blocks = []
n_flow_blocks = 2
mixture = VariationalMixture(
xdim=xdim,
hdim=hdim,
n_hidden=n_hidden,
n_classes=n_classes,
components=[
make_nf(shared_blocks, n_flow_blocks,
distrib.Normal(loc=torch.zeros(2), scale=torch.ones(2)))
for _ in range(n_classes)
],
)
# %%
f1 = AffineLUFlow(2)
# %%
f1.forward(X)
# %%
#blocks = sum(
# [[AffineLUFlow(2), PReLUFlow(2), BatchNormFlow(2)] for _ in range(5)] +
# [[AffineLUFlow(2)]],
#[])
blocks = sum(
[[StructuredAffineFlow(2),
PReLUFlow(2), BatchNormFlow(2)]
for _ in range(5)] + [[StructuredAffineFlow(2)]], [])
flow = NormalizingFlow(
*blocks,
base_dist=base_dist,
)
opt = optim.Adam(flow.parameters(), lr=1e-2)
# %%
count_parameters(flow)
# %%
n_epochs = 2000
bs = 512
idx_0 = np.logical_and(X0[:, 0] < 0, X0[:, 1] < 0)
colors[idx_0] = 0
idx_1 = np.logical_and(X0[:, 0] >= 0, X0[:, 1] < 0)
colors[idx_1] = 1
idx_2 = np.logical_and(X0[:, 0] >= 0, X0[:, 1] >= 0)
colors[idx_2] = 2
idx_3 = np.logical_and(X0[:, 0] < 0, X0[:, 1] >= 0)
colors[idx_3] = 3
# %%
plt.scatter(X0[:, 0], X0[:, 1], s=5, c=colors)
# %%
blocks = sum(
[[StructuredAffineFlow(2), PReLUFlow(2), BatchNormFlow(2)] for _ in range(20)] + [[StructuredAffineFlow(2)]],
[])
# %%
flow = NormalizingFlow(
*blocks,
base_dist=base_dist,
)
opt = optim.Adam(flow.parameters(), lr=1e-2)
# %%
count_parameters(flow)
# %%
n_epochs = 150
density = prev_density.squeeze() / np.exp(flow.log_abs_det_jacobian(torch.Tensor(cur_z), None).detach().squeeze())
return torch.Tensor(density)
# %%
def get_meshes(cur_z, density, grid_side=1000, dim=2):
mesh = cur_z.reshape([grid_side, grid_side, dim]).transpose(2, 0, 1)
xx = mesh[0]
yy = mesh[1]
zz = density.numpy().reshape([grid_side, grid_side])
return xx, yy, zz
# %%
blocks = [PReLUFlow(2)]
flow = NormalizingFlow(
*blocks,
base_dist=base_dist,
)
# %%
from mpl_toolkits.axes_grid1 import make_axes_locatable
# %%
x = np.linspace(-5, 5, 1000)
z = np.array(np.meshgrid(x, x)).transpose(1, 2, 0)
z = np.reshape(z, [z.shape[0] * z.shape[1], -1])
with torch.no_grad():