# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: py37
#     language: python
#     name: py37
# ---

# ## Dendrogram flow on KITTI Mask dataset

# %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

# ### KittiMask

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))

# ### Warm start

# Initialize the transition noise to Laplacian sparse noise

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")

# #### Dendrogram flow

# 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")



StarcoderdataJupyterScriptsDedupFiltered