train_cnf.py

import numpy as np
import matplotlib.pyplot as plt
import os
import glob
from PIL import Image
from functools import partial
import jax
from typing import Any, Callable, Sequence, Optional, NewType
from jax import lax, random, vmap, scipy, numpy as jnp
# from jax.experimental.ode import odeint
from models.ode import odeint
import flax
from flax.training import train_state
from flax import traverse_util
from flax.core import freeze, unfreeze
from flax import linen as nn
from flax import serialization
import optax
from sklearn.datasets import make_circles, make_moons, make_s_curve
from tqdm import tqdm


os.environ['CUDA_VISIBLE_DEVICES'] = '3'
# os.environ['TF_FORCE_UNIFIED_MEMORY'] = '1'
# os.environ['XLA_PYTHON_CLIENT_ALLOCATOR'] = 'platform'
# os.environ['XLA_PYTHON_CLIENT_PREALLOCATE'] = 'false'


class HyperNetwork(nn.Module):
    """Hyper-network allowing f(z(t), t) to change with time.

    Adapted from the Pytorch implementation at:
    https://github.com/rtqichen/torchdiffeq/blob/master/examples/cnf.py
    """
    in_out_dim: Any = 2
    hidden_dim: Any = 32
    width: Any = 64

    @nn.compact
    def __call__(self, t):
        # predict params
        blocksize = self.width * self.in_out_dim
        params = lax.expand_dims(t, (0, 1))
        params = nn.Dense(self.hidden_dim)(params)
        params = nn.tanh(params)
        params = nn.Dense(self.hidden_dim)(params)
        params = nn.tanh(params)
        params = nn.Dense(3 * blocksize + self.width)(params)

        # restructure
        params = lax.reshape(params, (3 * blocksize + self.width,))
        W = lax.reshape(params[:blocksize], (self.width, self.in_out_dim, 1))

        U = lax.reshape(params[blocksize:2 * blocksize], (self.width, 1, self.in_out_dim))

        G = lax.reshape(params[2 * blocksize:3 * blocksize], (self.width, 1, self.in_out_dim))
        U = U * nn.sigmoid(G)

        B = lax.expand_dims(params[3 * blocksize:], (1, 2))
        return W, B, U


class CNF(nn.Module):
    """Adapted from the Pytorch implementation at:
    https://github.com/rtqichen/torchdiffeq/blob/master/examples/cnf.py
    """
    in_out_dim: Any = 2
    hidden_dim: Any = 32
    width: Any = 64

    @nn.compact
    def __call__(self, t, states):
        z, logp_z = states[:, :2], states[:, 2:]
        W, B, U = HyperNetwork(self.in_out_dim, self.hidden_dim, self.width)(t)

        def dzdt(z):
            h = nn.tanh(vmap(jnp.matmul, (None, 0))(z, W) + B)
            return jnp.matmul(h, U).mean(0)

        # Naive version
        # def h(z):
        #     return nn.tanh(vmap(jnp.matmul, (None, 0))(z, W) + B)
        # dh_dz = jax.jacrev(h)(z)
        # matmul_1 = vmap(jnp.matmul, in_axes=(3, None), out_axes=3)
        # matmul_2 = vmap(matmul_1, in_axes=(4, None), out_axes=4)
        # u_dot_dh_dz = matmul_2(dh_dz, U).mean(0).sum(0)
        # dlogp_z_dt = -1.0 * jnp.trace(u_dot_dh_dz, 0, 0, 2)

        # More advanced.
        # u_dot_dh_dz = jax.jacrev(dzdt)(z).sum(0)
        # dlogp_z_dt = -1.0 * jnp.trace(u_dot_dh_dz, 0, 0, 2)

        # Final version
        dz_dt = dzdt(z)
        sum_dzdt = lambda z: dzdt(z).sum(0)
        df_dz = jax.jacrev(sum_dzdt)(z)
        dlogp_z_dt = -1.0 * jnp.trace(df_dz, 0, 0, 2)

        return lax.concatenate((dz_dt, lax.expand_dims(dlogp_z_dt, (1,))), 1)


class Neg_CNF(nn.Module):
    """Negative CNF for jax's odeint."""
    in_out_dim: Any = 2
    hidden_dim: Any = 32
    width: Any = 64

    @nn.compact
    def __call__(self, t, states):
        outputs = CNF(self.in_out_dim, self.hidden_dim, self.width)(-1.0 * t, states)

        return -1.0 * outputs


def get_batch_circles(num_samples):
    """Adapted from the Pytorch implementation at:
    https://github.com/rtqichen/torchdiffeq/blob/master/examples/cnf.py
    """
    points, _ = make_circles(n_samples=num_samples, noise=0.06, factor=0.5)
    x = jnp.array(points, dtype=jnp.float32)
    logp_diff_t1 = jnp.zeros((num_samples, 1), dtype=jnp.float32)

    return lax.concatenate((x, logp_diff_t1), 1)


def get_batch_moons(num_samples):
    points, _ = make_moons(n_samples=num_samples, noise=0.05)
    x = jnp.array(points, dtype=jnp.float32)
    logp_diff_t1 = jnp.zeros((num_samples, 1), dtype=jnp.float32)

    return lax.concatenate((x, logp_diff_t1), 1)


def get_batch_scurve(num_samples):
    points, _ = make_s_curve(n_samples=num_samples, noise=0.05, random_state=0)
    x1 = jnp.array(points, dtype=jnp.float32)[:, :1]
    x2 = jnp.array(points, dtype=jnp.float32)[:, 2:]
    x = lax.concatenate((x1, x2), 1)
    logp_diff_t1 = jnp.zeros((num_samples, 1), dtype=jnp.float32)

    return lax.concatenate((x, logp_diff_t1), 1)


def multivariate_normal(z):
    """
    Log probability of multivariate_normal.
    """
    mean = jnp.array([0., 0.])
    z_m = z - mean
    cov = jnp.array([[0.1, 0.], [0., 0.1]])
    logz = -jnp.log((2 * jnp.pi)) + -0.5 * jnp.log(jnp.linalg.det(cov)) + -0.5 * jnp.matmul(jnp.matmul(z_m.T, jnp.linalg.inv(cov)), z_m)
    return logz


def create_train_state(rng, learning_rate, in_out_dim, hidden_dim, width):
    """Creates initial 'TrainState'."""
    inputs = jnp.ones((1, 2))
    neg_cnf = Neg_CNF(in_out_dim, hidden_dim, width)
    params = neg_cnf.init(rng, jnp.array(10.), inputs)['params']
    set_params(params)
    tx = optax.adam(learning_rate)
    return train_state.TrainState.create(
        apply_fn=neg_cnf.apply, params=params, tx=tx
    )


def set_params(params):
    # Convert all value of Params to certain constant
    params = unfreeze(params)
    # Get flattened-key: value list.
    flat_params = {'/'.join(k): v for k, v in traverse_util.flatten_dict(params).items()}
    unflat_params = traverse_util.unflatten_dict({tuple(k.split('/')): 0.1 * jnp.ones_like(v) for k, v in flat_params.items()})
    new_params = freeze(unflat_params)
    test_x = jnp.array([[0., 1.], [2., 3.], [4., 5.]])
    test_log_p = jnp.zeros((3, 1))
    test_inputs = lax.concatenate((test_x, test_log_p), 1)
    Neg_CNF().apply({'params': new_params}, jnp.array(0.), test_inputs)


@partial(jax.jit, static_argnums=(2, 3, 4, 5, 6))
def train_step(state, batch, in_out_dim, hidden_dim, width, t0, t1):
    p_z0 = lambda x: scipy.stats.multivariate_normal.logpdf(x,
                                                            mean=jnp.array([0., 0.]),
                                                            cov=jnp.array([[0.1, 0.], [0., 0.1]]))
    def loss_fn(params):
        func = lambda states, t: Neg_CNF(in_out_dim, hidden_dim, width).apply({'params': params}, t, states)
        outputs, _ = odeint(
            func,
            batch,
            -1.0 * jnp.array([t1, t0]),
            atol=1e-5,
            rtol=1e-5
        )
        z_t, logp_diff_t = outputs[:, :, :2], outputs[:, :, 2:]
        z_t0, logp_diff_t0 = z_t[-1], logp_diff_t[-1]
        logp_x = p_z0(z_t0) - lax.squeeze(logp_diff_t0, dimensions=(1,))
        loss = -logp_x.mean(0)
        return loss
    grad_fn = jax.value_and_grad(loss_fn)
    loss, grads = grad_fn(state.params)
    state = state.apply_gradients(grads=grads)

    return state, loss


def train(learning_rate, n_iters, batch_size, in_out_dim, hidden_dim, width, t0, t1, visual, dataset):
    """Train the model."""
    rng = jax.random.PRNGKey(0)
    state = create_train_state(rng, learning_rate, in_out_dim, hidden_dim, width)
    if dataset == "circles":
        get_batch = lambda num_samples: get_batch_circles(num_samples)
    elif dataset == "moons":
        get_batch = lambda num_samples: get_batch_moons(num_samples)
    elif dataset == "scurve":
        get_batch = lambda num_samples: get_batch_scurve(num_samples)

    for itr in range(1, n_iters+1):
        batch = get_batch(batch_size)
        state, loss = train_step(state, batch, in_out_dim, hidden_dim, width, t0, t1)
        print("iter: %d, loss: %.2f" % (itr, loss))

    if visual is True:
        # Convert Params of Neg_CNF to CNF
        neg_params = state.params
        neg_params = unfreeze(neg_params)
        # Get flattened-key: value list.
        neg_flat_params = {'/'.join(k): v for k, v in traverse_util.flatten_dict(neg_params).items()}
        pos_flat_params = {key[6:]: jnp.array(np.array(neg_flat_params[key])) for key in list(neg_flat_params.keys())}
        pos_unflat_params = traverse_util.unflatten_dict({tuple(k.split('/')): v for k, v in pos_flat_params.items()})
        pos_params = freeze(pos_unflat_params)
        jax.profiler.save_device_memory_profile("memory.prof")
        output = viz(neg_params, pos_params, in_out_dim, hidden_dim, width, t0, t1, dataset)
        z_t_samples, z_t_density, logp_diff_t, viz_timesteps, target_sample, z_t1 = output
        create_plots(z_t_samples, z_t_density, logp_diff_t, t0, t1, viz_timesteps, target_sample, z_t1, dataset)


def solve_dynamics(dynamics_fn, initial_state, t):
    def f(initial_state, t):
        return odeint(dynamics_fn, initial_state, t, atol=1e-5, rtol=1e-5)
    return f(initial_state, t)


def viz(neg_params, pos_params, in_out_dim, hidden_dim, width, t0, t1, dataset):
    """Adapted from PyTorch """
    viz_samples = 30000
    viz_timesteps = 41
    if dataset == "circles":
        get_batch = lambda num_samples: get_batch_circles(num_samples)
    elif dataset == "moons":
        get_batch = lambda num_samples: get_batch_moons(num_samples)
    elif dataset == "scurve":
        get_batch = lambda num_samples: get_batch_scurve(num_samples)
    target_sample = get_batch(viz_samples)[:, :2]

    if not os.path.exists('results_%s/' % dataset):
        os.makedirs('results_%s/' % dataset)

    z_t0 = jnp.array(np.random.multivariate_normal(mean=np.array([0., 0.]),
                                                   cov=np.array([[0.1, 0.], [0., 0.1]]),
                                                   size=viz_samples))
    logp_diff_t0 = jnp.zeros((viz_samples, 1), dtype=jnp.float32)

    func_pos = lambda states, t: CNF(in_out_dim, hidden_dim, width).apply({'params': pos_params}, t, states)
    output = solve_dynamics(func_pos, lax.concatenate((z_t0, logp_diff_t0), 1), jnp.linspace(t0, t1, viz_timesteps))
    z_t_samples, _ = output[..., :2], output[..., 2:]

    # Generate evolution of density
    x = jnp.linspace(-1.5, 1.5, 100)
    y = jnp.linspace(-1.5, 1.5, 100)
    points = np.vstack(np.meshgrid(x, y)).reshape([2, -1]).T

    z_t1 = jnp.array(points, dtype=jnp.float32)
    logp_diff_t1 = jnp.zeros((z_t1.shape[0], 1), dtype=jnp.float32)
    func_neg = lambda states, t: Neg_CNF(in_out_dim, hidden_dim, width).apply({'params': neg_params}, t, states)
    output = solve_dynamics(func_neg, lax.concatenate((z_t1, logp_diff_t1), 1), -jnp.linspace(t1, t0, viz_timesteps))
    z_t_density, logp_diff_t = output[..., :2], output[..., 2:]

    return z_t_samples, z_t_density, logp_diff_t, viz_timesteps, target_sample, z_t1


def create_plots(z_t_samples, z_t_density, logp_diff_t, t0, t1, viz_timesteps, target_sample, z_t1, dataset):
    # Create plots for each timestep
    for (t, z_sample, z_density, logp_diff) in zip(
            tqdm(np.linspace(t0, t1, viz_timesteps)),
            z_t_samples, z_t_density, logp_diff_t
    ):
        fig = plt.figure(figsize=(12, 4), dpi=200)
        plt.tight_layout()
        plt.axis('off')
        plt.margins(0, 0)
        fig.suptitle(f'{t:.2f}s')

        ax1 = fig.add_subplot(1, 3, 1)
        ax1.set_title('Target')
        ax1.get_xaxis().set_ticks([])
        ax1.get_yaxis().set_ticks([])
        ax2 = fig.add_subplot(1, 3, 2)
        ax2.set_title('Samples')
        ax2.get_xaxis().set_ticks([])
        ax2.get_yaxis().set_ticks([])
        ax3 = fig.add_subplot(1, 3, 3)
        ax3.set_title('Log Probability')
        ax3.get_xaxis().set_ticks([])
        ax3.get_yaxis().set_ticks([])

        ax1.hist2d(*jnp.transpose(target_sample), bins=300, density=True,
                   range=[[-1.5, 1.5], [-1.5, 1.5]])

        ax2.hist2d(*jnp.transpose(z_sample), bins=300, density=True,
                   range=[[-1.5, 1.5], [-1.5, 1.5]])
        p_z0 = lambda x: scipy.stats.multivariate_normal.logpdf(x,
                                                                mean=jnp.array([0., 0.]),
                                                                cov=jnp.array([[0.1, 0.], [0., 0.1]]))
        logp = p_z0(z_density) - lax.squeeze(logp_diff, dimensions=(1,))
        ax3.tricontourf(*jnp.transpose(z_t1),
                        jnp.exp(logp), 200)

        plt.savefig(os.path.join('results_%s/' % dataset, f"cnf-viz-{int(t * 1000):05d}.jpg"),
                    pad_inches=0.2, bbox_inches='tight')
        plt.close()

    img, *imgs = [Image.open(f) for f in sorted(glob.glob(os.path.join('results_%s/' % dataset, f"cnf-viz-*.jpg")))]
    img.save(fp=os.path.join('results_%s/' % dataset, "cnf-viz.gif"), format='GIF', append_images=imgs,
             save_all=True, duration=250, loop=0)

    print('Saved visualization animation at {}'.format(os.path.join('results_%s/' % dataset, "cnf-viz.gif")))


if __name__ == '__main__':
    train(0.001, 1000, 512, 2, 32, 64, 0., 10., True, 'scurve')