Training Tutorial

[Davidmolin4]

Hello everybody!

I would like to know if there is any tutorial about how to train normalizing flows with this library. I've been looking but I haven't found any, I couldn't understand the FlowTrainer class either.

Greetings

David

[EddieCunningham]

Thanks for the interest in NuX! The flows are built to be as flexible as possible, so you could in theory train them using any of the training frameworks out there for JAX (optax, flax, haiku, etc.). You'd create and initialize the flow using the format in the readme example which is roughly

# 1. Create flow
flow = nux.Sequential(...., prior=...)

# 2. Initialize the flow with data
z, log_px = flow(initialization_data, params=None, rng_key=initialization_key)

# 3. Retrieve the initialized parameters
params = flow.get_params()

# 4. Give your JAX framework the parameters
# (Depends on framework.  See [here](https://github.com/Information-Fusion-Lab-Umass/NuX/blob/master/examples/haiku_vae.py#L83) for a Haiku example)

# 5. Use the flow in your code
z, log_px = flow(x, params=params, rng_key=rng_key)

The FlowTrainer class is another way to train a model. You can take a look at this for an example of how to use it, but to be honest I think training with some of the larger frameworks might be the best way to go because they are more developed. However if you're still interested in how to train flows using the class in this library and want to learn more, I'd be happy to explain more about it.

[EddieCunningham]

Hmm, I'm not sure what the issue could be. When I run your code the loss goes down to 1.43 after 300 steps and stays there. Do you have the latest version of NuX? Older versions have had bugs in the spline class.

Also, you can significantly speed up this example by jitting the gradient function! You would do this by changing your training loop to

jitted_grad = jax.jit(jax.grad(loss))
for step in tqdm(range(epochs)):
  if step%100 == 0:
    print(step,'%.2f' % (loss(params,x, rng_key)))
  grads = jitted_grad(params, x, rng_key)
  updates, opt_state = optimizer.update(grads, opt_state)
  params = optax.apply_updates(params, updates)

[Davidmolin4]

Thanks again Eddie

I installed nux in colab through the following command

!git clone --recursive https://github.com/Information-Fusion-Lab-Umass/NuX.git

But I can't tell which version has been installed

[EddieCunningham]

Can you share a copy of the colab notebook? I'm not able to reproduce the error.

[Davidmolin4]

I discovered the error.

I was using the following data.

# No gaussian data
rng_key = jax.random.PRNGKey(0)
batch_size, dim = 10000, 1
x = jax.random.normal(rng_key, (batch_size, dim))
x = jnp.concatenate((x, jax.random.normal(rng_key, (batch_size, dim))+10))

With the default parameter bounds for the spline. Changing the bounds according the data solves the issue.

I wrote you in the question another data. My mistake.

[EddieCunningham]

No problem, I'm glad you found the issue