Speed up loop...

[jecampagne]

Hi, I'm pretty sure that this question will be 5' solved by experts.

Here is a working snippet

import numpy as np

from jax import random as rnd
# JAX numpy implementation
import jax.numpy as jnp
# JAX Scipy implementation
from jax.scipy.linalg import cholesky
from jax.scipy.stats import multivariate_normal
from jax import vmap  # vectorization
from jax import jit   # code optimization
from jax import value_and_grad  # auto-diff gradient 
from jax.example_libraries.optimizers import adam, exponential_decay  # Adam minimizer

def transform_to_covariance_matrix(sq_mat):
    """ transform a squarred matrix into a covariance matrix"""
    U = jnp.triu(sq_mat)
    U_T = jnp.transpose(U)
    return jnp.dot(U_T, U)

##True Mean and Covariant mtx
key, key_mu,key_cov,key_data = rnd.split(rnd.PRNGKey(42),4)
true_mu  = rnd.normal(key_mu, shape=(5,))
true_cov = transform_to_covariance_matrix(rnd.normal(key_cov, shape=(5, 5)))  # let's make a 5x5 covariance matrix

##SOme data from the true distrib
data = rnd.multivariate_normal(key_data,true_mu,true_cov,(10000,))
# Let me split into batches
data_batched = data.reshape(10,-1,5)

### The -log-likelihood
def loglike(params, data):
    mu, untransformed_cov = params
    cov = transform_to_covariance_matrix(untransformed_cov)
    def logpdf_func(obs):
        """logpdf of multivariate normal for one observation"""
        return multivariate_normal.logpdf(obs,mu, cov)
    logp = vmap(logpdf_func)(data)
    return -np.sum(logp)

# JAX value & gradient of log-likelihood
dloglike = value_and_grad(loglike)

# Get the optimizer utils
init, update, get_params = adam(step_size=5e-5)

# Optimize the code
get_params = jit(get_params)
update = jit(update)
dloglike = jit(dloglike)
loglike = jit(loglike)

# init parameters
key, key_mu, key_cov = rnd.split(key,3)
mu = rnd.normal(key_mu, shape=(5,))
untransformed_cov = rnd.normal(key_cov, shape=(5, 5))
params = mu, untransformed_cov
state = init(params)
loss_rate = []

#Perform the Gradient Descent minimization
for i in range(3_000):                      # epoch loop
    params = get_params(state)
    for ib in range(10):                       #batch loop
        obs = data_batched[ib]
        val,g = dloglike(params, obs)
        state = update(i, g, state)
        loss_rate.append(val)

# Get the optimized Mean & Cov mtx
mu_opt, untransformed_cov_opt = get_params(state)
cov_opt = transform_to_covariance_matrix(untransformed_cov_opt)

But it is quite slow few minutes on a CPU, and I get the double for-loops are horrible.
How I can speed up this code (nb. keeping the loss eval is here to plot diagnostics)

Thanks

[jecampagne]

Well I have managed to cook this code

def do_sgd(params,data_batched, n_epochs):
    n_batch = data_batched.shape[0]
    loss=jnp.zeros(n_epochs*n_batch)
    state = init(params)
    
    def body(i, carry):
        
        idx, loss, state = carry
        obs = data_batched[i]
        params = get_params(state)
        val,g = dloglike(params, obs)
        loss = loss.at[idx+i].set(val)
        state = update(i, g, state)
        
        carry = idx, loss, state
        
        return carry

    for epoch in range(n_epochs):
        idx = epoch*n_batch
        _,loss,state = lax.fori_loop(0, n_batch, body, (idx,loss,state))

    return loss, state

key, key_mu, key_cov = rnd.split(key,3)
mu = rnd.normal(key_mu, shape=(5,))
untransformed_cov = rnd.normal(key_cov, shape=(5, 5))
params = mu, untransformed_cov

loss, state_opt =  do_sgd(params,data_batched,5000)

It's working but I do not know if its the most efficient...

[YouJiacheng]

lax.fori_loop works? But I think there is a sequential dependency here.
Maybe you just need to combine get_params&value_and_grad&update in one function and jit it?

[soraros]

You could do something like this

def main(state, i):
  params = get_params(state)
  def f(state, obs):
    val, g = value_and_grad(loglike)(params, obs)
    state = update(i, g, state)
    return state, val
  return lax.scan(f, state, data_batched)

state, loss = jit(partial(lax.scan, main))(state, jnp.arange(5000))
loss = loss.ravel()

[jecampagne]

Great, it works perfectly. Thanks