SAve an object with dill... failed

[jecampagne]

Hi,
Well I try to save the current state of an object...

First let me gives a working case.

The imports

import jax
import jaxopt
import jax.numpy as jnp
import jax.scipy as jsc
from jax import vmap, jit
jax.config.update("jax_enable_x64", True)

import os
import dill

Now some I/O helpers:

def load_pkl_file(folder_name, file_name):
    '''
    Given a folder name and a file name, we will load the Python class.
    For example, a full GP module

    :param: folder_name (str) - the name of the folder

    :param: file_name (str) - name of the file

    :return: module (Python class) : complete module or it can be EMCEE full module
    '''


    with open(folder_name + '/' + file_name + '.pkl', 'rb') as f:
        module = dill.load(f)

    return module

def store_pkl_file(module, folder_name, file_name):
    '''
    Given a trained GP (module), we will save it given
    a folder name and a file name

    :param: module (python class) - for example, the GP module

    :param: folder_name (str) - the name of the folder

    :param: file_name (str) - name of the file
    '''


    # create the folder if it does not exist
    if not os.path.exists(folder_name):
        os.makedirs(folder_name)


    # store the module using dill
    with open(folder_name + '/' + file_name + '.pkl', 'wb') as f:
        dill.dump(module, f)

A simple class:

class A():
    def __init__(self,a,b):
        self.a = a
        self.b = b
        self.res = None

    def set_res(self,x):
        self.res = x
        
    def f(self):
        return self.a * self.b

Then, playing with it is fine

obja = A(10., 2.)

res = obja.f()

obja.set_res(res)

print(obja.a, obja.b, obja.res)

store_pkl_file(obja,'.','test')

newa = load_pkl_file('.','test')

print(newa.a, newa.b, newa.res)

This gives

10.0 2.0 20.0
10.0 2.0 20.0

and I can change a & b to use the f-func:

newa.a=30.
newa.b=3.
newa.f()

which leads to 90.0 as expected.

So far so good, dillis working with a regular Python class, ie. no problem.
Let me just change a little bit the A::f function

    def f(self):
        return self.a * jnp.exp(self.b)

it works too and we see that newa.resis an array

while

newa.a=30.
newa.b=3.
newa.f()

gives

DeviceArray(602.5661077, dtype=float64, weak_type=True)

Well, now come the real use case which is not working:

The class...

class GPEmu():
    """
        y(x) = \Phi(x).\beta + GP(0,K(x,x^\prime; \alpha)
        \alpha = (amplitude, \{\ell \}_i<d)  with d dimension of x
        K: GP kernel as RBF

        Determination of \beta & \alpha based on: 
        Gaussian process class
        Using C. E. Rasmussen & C. K. I. Williams, 
        Gaussian Processes for Machine Learning, the MIT Press, 2006 

        2.7 Incorporating Explicit Basis Functions
        Eqs. 2.39-2.41 (mean only)
        Eq. 2.43 Marginal Likelihood
        Hypothesis: beta prior=N(0,\lambda_cap Id) => simplififications
        
    """

    def __init__(self,
            kernel: Callable[[jnp.ndarray, 
                              jnp.ndarray, 
                              Dict[str, jnp.ndarray], 
                              float, float],jnp.ndarray] = kernel_RBF,
            var: float = 1e-5,
            order: int = 2,
            lambda_cap: float = 1.0,
            l_min: float = -5.0,
            l_max: float =  5.0,
            a_min: float =  0.0,
            a_max: float = 10.0,
            jitter: float = 1e-10,
            n_restart: int = 5,
            x_trans: bool = True,
            y_trans: bool = False,
            use_mean: bool = False) -> None:
        

        # order of the poynomial regression
        # we support only second order here
        self.order = order
        if self.order > 2:
            msg = 'At the moment, we support only order = 1 and order = 2'
            raise RuntimeError(msg)
        

        
        self.kernel = kernel  # Gaussian Process kernel
                        
        self.var = var        # "noise" level in principe can be set to very low level
        self.jitter = jitter  # to tweek choloesky algo add a small value on the diagonal if necessary
        self.order= order     # order of polynomial approx Phi(X)
        self.lambda_cap = lambda_cap # scale of polynomial coeff priors
        self.l_min = l_min # lower bound on kernel length (log scale)
        self.l_max = l_max # upper bound on kernel length (log scale)
        self.a_min = a_min # lower bound on kernel scale/amplitude (log scale)
        self.a_max = a_max # upper bound on kernel scale/amplitude (log scale)
        
        self.n_restart = n_restart #number of trials for kernel optimization in case of un-sucess

        # choose to make transformation
        self.x_trans = x_trans
        self.y_trans = y_trans

        # if we want to centre the output on zero
        self.use_mean = use_mean

        
        # init to None
        self.Init()
        
            
    def Init(self)->None:
        """
        Fnt used to reset internal data
        """
        self.theta   = None    # cosmo \Theta_i i<N_train
        self.y       = None    # f(\Theta_i)
        self.x_train = None    # Training input ( Ntrain x d ) with d: number of elements per train spl
        self.y_train = None    # Training target (Ntrain x 1)
        self.phi     = None    # Phi(X_train)
        self.nbasis  = None    # nber of elements Phi(X)  = 1+ order*d
        self.n_kern_params = None #  nber of Cosmo params + 1 = nber of K hyper-param

        # after kernel optimization (saved for prediction)
        self.kernel_hat  = None  # kernel parameters
        self.beta_hat    = None  # beta
        self.kinv_XX_res = None  # K(Xtrain, Xtrain)^{-1} * (y_train - Phi(X_train) * beta_hat)        

    #setter necessary as jit function cannot set self.<variable> (leaks) (eg. post_training)
    #Todo: think to use  @property
    def set_kernel_hat(self, x):
        self.kernel_hat = x

    def set_beta_hat(self, x):
        self.beta_hat = x
    
    def set_kinv_XX_res(self, x):
        self.kinv_XX_res = x
    
        

    def mean_Theta_y(self,theta: jnp.ndarray, y: jnp.ndarray)->None:
        """
        Compute mean_theta and center or not output
        """
        self.theta  = theta       
        self.y      = y           # f(\Theta_i)

        self.d = self.theta.shape[1]   # the dimension of the problem
        self.n_trains = self.theta.shape[0] # the number of training point
        msg = 'The number of training points is smaller than the dimension of the problem. Reshape your array!'
        # the number of training points is greater than the number of dimension
        assert self.n_trains > self.d, msg

        # compute mean of training set
        self.mean_theta = jnp.mean(theta, axis=0)
        # centre the input on zero
        self.theta = theta - self.mean_theta
        
        # if we want to centre the output on zero
        if self.use_mean:
            self.mean_function = jnp.mean(y)
        else:
            self.mean_function = jnp.zeros(1)

        # the output is of size (ntrain x 1)
        self.y = y.reshape(self.n_trains, 1) - self.mean_function


        
    def do_transformation(self) -> None:
        '''
        Perform all transformations on input theta, y to get x_train, y_train
        '''
        # we transform both x and y if specified
        if (self.x_trans and self.y_trans):
            self.transform = transformation(self.theta, self.y)
            self.x_train = self.transform.x_transform()
            self.y_train = self.transform.y_transform()

        # we transform x only if specified
        elif self.x_trans:
            print("do_transformation: x_trans only")
            self.transform = transformation(self.theta, self.y)
            self.x_train = self.transform.x_transform()
            self.y_train = self.y

        # we keep the inputs and outputs (original basis)
        else:
            self.x_train = self.theta
            self.y_train = self.y
            
        #
        self.y_train = self.y_train.squeeze()   # Todo: a voir si ensuite ajouter un indice n'est pas superflu

    def compute_basis(self, X: jnp.ndarray) -> jnp.ndarray:
        """
        Compute the input basis functions 
        .. math::
                \Phi = [1, \theta_1,\dots,\theta_p, \theta_1^2,\dots,\theta_p^2, \dots, \theta_p^d]
        """

        dummy_phi = [X**i for i in jnp.arange(1, self.order + 1)]
        phi = jnp.concatenate(dummy_phi, axis=1)
        phi = jnp.c_[np.ones((X.shape[0], 1)), phi]
        return phi

    
    def basis_prior(self, par: Dict[str,Number])->Dict[str,jnp.array]:
        """
        beta ~ N(mu,C) with mu=0, C=lambda_cap * Id
        """
        
        n = par["nbasis"]
        lambda_cap = par["lambda_cap"]

        assert lambda_cap>0, f"lambda_cap should be >0 got {lambda_cap}" 

        beta = numpyro.sample("beta",
                              dist.MultivariateNormal(
                                  loc=jnp.zeros((n,)), 
                                    covariance_matrix=lambda_cap * jnp.identity(n), 
                                    precision_matrix = 1./lambda_cap * jnp.identity(n),
                              ),
                            )
        return {"beta": beta[:,jnp.newaxis]}
    
    
    def kernel_prior(self)->Dict[str,jnp.array]:
        """
        Kernel hyper-parameter priors
        k_scale (aka amplitude): 1 parameter
        k_length: as many as comosmological parameter
        """
        
        k_scale = numpyro.sample("k_scale",dist.Uniform(jnp.exp(self.a_min),jnp.exp(self.a_max)))
        tmp = jnp.repeat(jnp.array([[0., 1.0]]), self.n_lengthes, axis=0)
        k_length = numpyro.sample("k_length", dist.LogNormal(tmp[:,0],tmp[:,1]))

        return {"k_scale":k_scale, "k_length":k_length}


    def prepare_training(self)->None:
        """
        Before training/optimization
        """
        self.phi    = self.compute_basis(self.x_train)   # (N x m) with m= order * n_cosmo +1)
        self.nbasis = self.phi.shape[1]       # m
        self.par_basis_prior = {"nbasis":self.nbasis, "lambda_cap": self.lambda_cap}

        self.n_lengthes = self.x_train.shape[1]
        self.n_kern_params = self.n_lengthes + 1
        
        print(f"prepare_training: {self.par_basis_prior}, {self.n_kern_params}")

    
    
    #@partial(jit, static_argnums=(0,))
    def marginal_likelihood(self, params:jnp.array) -> float:
        """
        Marginal Likelihood with hypothesis of Gaussian Prior for beta
        """
        
        k_scale =  params[0]
        k_length = params[1:]
        kernel_params= {"k_scale":k_scale, "k_length":k_length}
        
        # beta prior = N(0, lambda_cap * Id)
        # basis_cov = lambda_cap \Phi(X_train) \Phi(X_train)^T  (NxN matrix)
        # basis_cov can be computed once but NxN matrix -> save storage = compute each time
        
        basis_cov = self.lambda_cap * (self.phi @ self.phi.T)
        
        # compute GP K(X,X)
        KGP = self.kernel(self.x_train, self.x_train, 
                        params=kernel_params,
                        noise=self.var, 
                        jitter=self.jitter)
        
        # Kernel tot 
        K_XX = KGP + basis_cov
        
        # the beta prior mean is 0
        y_residuals = self.y_train[:,jnp.newaxis]
        
        chol_XX = jsc.linalg.cholesky(K_XX, lower=True)
        v  = jsc.linalg.solve_triangular(chol_XX, y_residuals, lower=True)
        
        mlp  = -0.5 * jnp.dot(v.T,v)                       # chi2
        mlp += -0.5 * jnp.sum(jnp.log(jnp.diag(chol_XX)))  # complexity term (determinant)
        mlp += -0.5 * self.n_trains * jnp.log(2.*jnp.pi)   # cte not really necessary in the min
        
        return mlp.squeeze()
        
    ## ERROR due to ScipyBoundedMinimize @partial(jit, static_argnums=(0,)) -> Issue JaxOpt
    def optimize(self, rng_key:jnp.array) -> bool:
        
        print_summary = False
        #minimizer
        jscMin=jaxopt.ScipyBoundedMinimize(fun=lambda x : -self.marginal_likelihood(x),
                                   method='L-BFGS-B', options={'disp':print_summary})
        
        bnd_low = [self.a_min]+[self.l_min]*self.n_lengthes
        bnd_high =[self.a_max]+[self.l_max]*self.n_lengthes
        bnd_low = jnp.exp(jnp.array(bnd_low))
        bnd_high = jnp.exp(jnp.array(bnd_high))
#        print(f"bnd_low: {bnd_low}")
#        print(f"bnd_high: {bnd_high}")
        
        # initial guess
        rng_key, new_key = jax.random.split(rng_key)
        init_par = trace(seed(self.kernel_prior,new_key)).get_trace()
        init_param = jnp.hstack((init_par['k_scale']['value'], init_par['k_length']['value']))
        
#        print(f'init_param: {init_param}')
        
        # do minimization
        res = jscMin.run(init_param,bounds=(bnd_low, bnd_high))

        # optimized parameters and status
        param_hat, state = res
        
#        print("state:",state)
        
        return state.success, {'k_scale':param_hat[0],'k_length':param_hat[1:]}

    #test: to use jit while cannot be done in opimize due to ScipyBoundedMinimize
    #@partial(jit, static_argnums=(0,))
    def post_training(self)->None:
        
        # compute some quanties for predicton
        k_XX = self.kernel(self.x_train, self.x_train, params=self.kernel_hat, 
                         noise=self.var,jitter=0.0)  #aka Ky
                
        chol_XX = jax.lax.linalg.cholesky(k_XX)    # K_XX = Lxx Lxx^T
        
        
        
        # beta_hat = W_\beta^{-1} [\Phi^T K_y^{-1} y + C^{-1}\mu]      with \mu=0 C=\lambda_{cap} Id
        # W_\beta = C^{-1} + \Phi^T K_y^{-1} \Phi

        #kinv_XX_y = K_y^{-1} y   => kinv_XX_y=Lxx^T\(Lxx\y)  with  
        kinv_XX_y = jax.lax.linalg.triangular_solve(
            chol_XX.T, 
            jax.lax.linalg.triangular_solve(chol_XX, self.y_train[:,jnp.newaxis],
                                            lower=True,  left_side=True),
            left_side=True)

        v1 = self.phi.T @ kinv_XX_y
        v2 = jax.lax.linalg.triangular_solve(chol_XX, self.phi, lower=True, left_side=True)

        Wb = v2.T @ v2 + 1./self.lambda_cap * jnp.identity(self.nbasis)

        chol_Wb = jax.lax.linalg.cholesky(Wb)  # Wb = Lb Lb^T

        # beta_hat = W_\beta^{-1} v1 => beta_hat = Lb^T\(Lb\y)
        beta_hat = jax.lax.linalg.triangular_solve(
            chol_Wb.T, 
            jax.lax.linalg.triangular_solve(chol_Wb, v1, lower=True,  left_side=True),
            left_side=True)

        # y_residuals = y_train - \Phi(X_train)\hat{\beta}
        y_residuals = self.y_train[:,jnp.newaxis] - self.phi @ beta_hat

        #kinv_XX_res = \hat{K_y}^{-1} y_residuals
        kinv_XX_res = jax.lax.linalg.triangular_solve(
                    chol_XX.T, jax.lax.linalg.triangular_solve(chol_XX, y_residuals, lower=True,  
                                                               left_side=True),
                    left_side=True)
        return beta_hat, kinv_XX_res
    
    
    #@partial(jit, static_argnums=(0,))
    def simple_predict(self, X_new: jnp.ndarray,) -> jnp.ndarray:
        """
        Make prediction after optimization
        y_\ast = \Phi_\ast \hat{\beta} 
                + \hat{k_\ast}^T \hat{K_y}^{-1} (y_train - \Phi(X_train)\hat{\beta})        
        """

        X_new = X_new if X_new.ndim > 1 else X_new[:, jnp.newaxis]  # possibly single test
        phi_star = self.compute_basis(X_new)   # (N* x m)
        k_pX = self.kernel(X_new, self.x_train, self.kernel_hat, noise=0.0, jitter=self.jitter)

        y_new = phi_star @ self.beta_hat + jnp.matmul(k_pX, self.kinv_XX_res)  
    
        return y_new.squeeze()

gp_model = GPEmu(kernel=kernel_RBF,
                         var=st.var,
                         order=st.order,
                         lambda_cap= st.gf_args['lambda_cap'],
                         l_min=st.l_min,
                         l_max=st.l_max,
                         a_min=st.a_min,
                         a_max=st.a_max,
                         jitter=st.jitter,
                         x_trans=st.x_trans,
                         y_trans=st.gf_args['y_trans'],
                         use_mean=st.use_mean)

I use it now like that

def Job(theta: np.ndarray,
        y: np.ndarray,
        gp_model: GPEmu, 
        gp_args: Dict,
        n_restart: int = 5):
              
        gp_model.Init()
        gp_model.mean_Theta_y(theta,y)

        #theta,y => x_train, y_train (using x_trans & y_trans)
        print("do_transformation")
        gp_model.do_transformation()  
        
        # prepare training
        print("prepare_training")
        gp_model.prepare_training()
        
        #do training (optimization of Kernel hyper-parameters
        rng_key = jax.random.PRNGKey(42)
            
        #Todo : switch to jax.lax.while_loop(cond_fun,body,val)
        print("optimize...")
        rc = False
        n_iter=0
        while  (not rc) and (n_iter<n_restart) :
            rng_key, new_key = jax.random.split(rng_key)
            rc, kernel_para = gp_model.optimize(new_key)
            print(f"try: {n_iter}, rc: {rc}")
            n_iter += 1

        #compute qties for predictions (care: jit fnt cannot use self.<var> = ...) (leaks)
        print("post training...")
        gp_model.set_kernel_hat(kernel_para)
        beta_hat, kinv_XX_res = gp_model.post_training()
        gp_model.set_beta_hat(beta_hat)
        gp_model.set_kinv_XX_res(kinv_XX_res)
        
        #Store gp model: Todo do a cleaning before
        print("store...")
        #store_pkl_file(gp_model, gp_args["folder_name"],gp_args["file_name"])

        #end
        print("Done")
        return gp_model

gf_job = Job(theta=theta, y=y, 
             gp_model=gp_model,
            gp_args=arg_cur_gp,             
            n_restart= st.n_restart
            )

While the result of the algorithm is exactly what I am expecting, Here is the crash/error trace back when I try to store:

store_pkl_file(gf_job, '.','tmp')

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Input In [104], in <cell line: 1>()
----> 1 store_pkl_file(gf_job, '.','tmp')

File helper.py:86, in store_pkl_file(module, folder_name, file_name)
     84 # store the module using dill
     85 with open(folder_name + '/' + file_name + '.pkl', 'wb') as f:
---> 86     dill.dump(module, f)

File ~/.local/lib/python3.8/site-packages/dill/_dill.py:276, in dump(obj, file, protocol, byref, fmode, recurse, **kwds)
    274 _kwds = kwds.copy()
    275 _kwds.update(dict(byref=byref, fmode=fmode, recurse=recurse))
--> 276 Pickler(file, protocol, **_kwds).dump(obj)
    277 return

File ~/.local/lib/python3.8/site-packages/dill/_dill.py:498, in Pickler.dump(self, obj)
    496     raise PicklingError(msg)
    497 else:
--> 498     StockPickler.dump(self, obj)
    499 stack.clear()  # clear record of 'recursion-sensitive' pickled objects
    500 return

File /python3.8/pickle.py:487, in _Pickler.dump(self, obj)
    485 if self.proto >= 4:
    486     self.framer.start_framing()
--> 487 self.save(obj)
    488 self.write(STOP)
    489 self.framer.end_framing()

File /python3.8/pickle.py:603, in _Pickler.save(self, obj, save_persistent_id)
    599     raise PicklingError("Tuple returned by %s must have "
    600                         "two to six elements" % reduce)
    602 # Save the reduce() output and finally memoize the object
--> 603 self.save_reduce(obj=obj, *rv)

File /python3.8/pickle.py:717, in _Pickler.save_reduce(self, func, args, state, listitems, dictitems, state_setter, obj)
    715 if state is not None:
    716     if state_setter is None:
--> 717         save(state)
    718         write(BUILD)
    719     else:
    720         # If a state_setter is specified, call it instead of load_build
    721         # to update obj's with its previous state.
    722         # First, push state_setter and its tuple of expected arguments
    723         # (obj, state) onto the stack.

File /python3.8/pickle.py:560, in _Pickler.save(self, obj, save_persistent_id)
    558 f = self.dispatch.get(t)
    559 if f is not None:
--> 560     f(self, obj)  # Call unbound method with explicit self
    561     return
    563 # Check private dispatch table if any, or else
    564 # copyreg.dispatch_table

File ~/.local/lib/python3.8/site-packages/dill/_dill.py:990, in save_module_dict(pickler, obj)
    987     if is_dill(pickler, child=False) and pickler._session:
    988         # we only care about session the first pass thru
    989         pickler._session = False
--> 990     StockPickler.save_dict(pickler, obj)
    991     log.info("# D2")
    992 return

File /python3.8/pickle.py:971, in _Pickler.save_dict(self, obj)
    968     self.write(MARK + DICT)
    970 self.memoize(obj)
--> 971 self._batch_setitems(obj.items())

File /python3.8/pickle.py:997, in _Pickler._batch_setitems(self, items)
    995     for k, v in tmp:
    996         save(k)
--> 997         save(v)
    998     write(SETITEMS)
    999 elif n:

File /python3.8/pickle.py:603, in _Pickler.save(self, obj, save_persistent_id)
    599     raise PicklingError("Tuple returned by %s must have "
    600                         "two to six elements" % reduce)
    602 # Save the reduce() output and finally memoize the object
--> 603 self.save_reduce(obj=obj, *rv)

File /python3.8/pickle.py:687, in _Pickler.save_reduce(self, func, args, state, listitems, dictitems, state_setter, obj)
    684     raise PicklingError(
    685         "args[0] from __newobj__ args has the wrong class")
    686 args = args[1:]
--> 687 save(cls)
    688 save(args)
    689 write(NEWOBJ)

File /python3.8/pickle.py:560, in _Pickler.save(self, obj, save_persistent_id)
    558 f = self.dispatch.get(t)
    559 if f is not None:
--> 560     f(self, obj)  # Call unbound method with explicit self
    561     return
    563 # Check private dispatch table if any, or else
    564 # copyreg.dispatch_table

File ~/.local/lib/python3.8/site-packages/dill/_dill.py:1398, in save_type(pickler, obj)
   1396     log.info("# T6")
   1397     return
-> 1398 elif obj.__module__ == '__main__':
   1399     if issubclass(type(obj), type):
   1400     #   try: # used when pickling the class as code (or the interpreter)
   1401         if is_dill(pickler, child=True) and not pickler._byref:
   1402             # thanks to Tom Stepleton pointing out pickler._session unneeded

AttributeError: __module__

Sorry, I know that the above snippet is long and is not a working snippet, but to make it fine to play with it, is not so simple.

Does someone has an idea of what is going on?
Or, how do you proceed to save class state in JAX. (Notice that none of the GPemu class functions is jitted).
Thanks

[YouJiacheng]

a guess: some jnp/jsp function is jitted.
BTW, you can change the code where the error was raised, to print which obj doesnt have __module__ attr and locate the problem.