<jupyter_start><jupyter_text>Dendrogram flow on KITTI Mask dataset<jupyter_code>%load_ext autoreload
%autoreload 2
import ltcl
from ltcl.modules.trainer import Dendrogram
from ltcl.modules.dendrogram import AfflineDendrogramFlow
import torch
from torch.utils.data import DataLoader, random_split
import pytorch_lightning as pl
from pytorch_lightning.callbacks.early_stopping import EarlyStopping

import seaborn as sns
import matplotlib.pyplot as plt
%matplotlib inline
use_cuda = torch.cuda.is_available()
device = torch.device("cuda:3" if use_cuda else "cpu")
torch.backends.cudnn.benchmark = True<jupyter_output><empty_output><jupyter_text>KittiMask<jupyter_code>from ltcl.datasets.kitti import KittiMasks
import numpy as np
import numpy.random as random_state
dataset = KittiMasks("/home/cmu_wyao/projects/slow_disentanglement/data/kitti/")
train_data, val_data = random_split(dataset, [80000, 2506])
train_loader = DataLoader(train_data, batch_size=16, shuffle=True)
val_loader   = DataLoader(train_data, batch_size=128, shuffle=False)
batch = next(iter(train_loader))<jupyter_output><empty_output><jupyter_text>Warm start Initialize the transition noise to Laplacian sparse noise<jupyter_code>flow = AfflineDendrogramFlow(input_dims = (1, 64, 64),
                             width = 128,
                             depth = 32,
                             n_levels = 3,
                             lags = 1)
flow = flow.to(device)
spline_optimizer = torch.optim.Adam(flow.spline.parameters(), 
                                    lr=1e-3,
                                    weight_decay=0.0)

# Warm-start the prior to standard normal dist
for step in range(5000):
    latent_size = 64*64*1
    y = torch.distributions.laplace.Laplace(0,1).rsample((32, latent_size))
    dataset = y.to(device)
    spline_optimizer.zero_grad()
    z, logabsdet = flow.spline(dataset)
    logp = flow.base_dist.log_prob(z).sum(dim=1) + logabsdet
    loss = -torch.mean(logp)
    loss.backward(retain_graph=True)
    spline_optimizer.step()
sns.distplot(flow.sample_noise(28).detach().cpu().numpy()[:,150], hist=False, kde=True,
             bins=None,
             hist_kws={'edgecolor':'black'},
             kde_kws={'linewidth': 2}, label="prior");
torch.save(flow.state_dict(), "/home/cmu_wyao/projects/dendrogram_init.pth")<jupyter_output><empty_output><jupyter_text>Dendrogram flow<jupyter_code># Callbacks
early_stop_callback = EarlyStopping(
    monitor='val_loss',
    min_delta=0.00,
    patience=3,
    verbose=False,
    mode='min')
model = Dendrogram(input_dims = (1, 64, 64),
                   width = 64,
                   depth = 16,
                   n_levels = 3,
                   lags = 1,
                   lr = 1e-4) 
model = Dendrogram.load_from_checkpoint("/home/cmu_wyao/checkpoints/dendrogram_0.000000/lightning_logs/version_4/checkpoints/epoch=0-step=4999.ckpt", 
                                        input_dims = (1, 64, 64),
                   width = 128,
                   depth = 32,
                   n_levels = 3,
                   lags = 1,
                   lr = 1e-4)
trainer = pl.Trainer(default_root_dir = "/home/cmu_wyao/checkpoints/dendrogram",
                     gpus=2, 
                     max_epochs=50,
                     callbacks=[early_stop_callback])

# Train the model
trainer.fit(model, train_loader, val_loader)
sns.distplot(model.flow.sample_noise(10).detach().cpu().numpy()[:,70], hist=False, kde=True,
             bins=None,
             hist_kws={'edgecolor':'black'},
             kde_kws={'linewidth': 2}, label="prior");
model.flow
first_sample, second_sample, latents1, latents2 = batch
log_probs = model.flow.log_prob(second_sample, first_sample)
es, z, sum_log_abs_det_jacobians = model.flow(second_sample, first_sample)
x,_ = model.flow.inverse(z, first_sample)
x = x.detach().cpu().numpy()
plt.imshow(first_sample[1,0,0])
plt.imshow(x[1,0])
import seaborn as sns
sns.heatmap(C, cmap="YlGnBu")<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11427856964241061, 'fraction_numerical': 0.0277569392348087, 'mean_word_length': 2.9840637450199203, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '550740', 'license': 'MIT', 'n_tokens_mistral': 1470, 'n_tokens_neox': 1360, 'n_words': 235, 'path': 'tests/test_dendrogram.ipynb', 'repo_name': 'anonymous-authors-iclr2022-481/ltcl', 'size': '3999'}