Classifier class pb with JAX using a Python class

[jecampagne]

Hi,

I face a number of problems with a quite simple classifier application that I would like to code in JAX to show to some students...

import time

import jax
import jax.numpy as jnp
from jax import random
from jax import grad, jit, vmap, value_and_grad
from jax.example_libraries import optimizers

from jax.config import config
config.update("jax_enable_x64", True)


import numpy as np

import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.pyplot import cm
plt.style.use('seaborn-white')
mpl.rc('image', cmap='jet')
mpl.rcParams['font.size'] = 18

print("DEBUG Jax device:",jax.devices())

from sklearn import datasets
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split


from jax.example_libraries import stax
from jax.example_libraries.stax import (BatchNorm, Conv, Dense, Flatten,
                                   Relu, LogSoftmax)

from torch.utils.data import Dataset,  DataLoader



def one_hot(x, k, dtype=np.float32):
    """Create a one-hot encoding of x of size k """
    return jnp.array(x[:, None] == jnp.arange(k), dtype)

n_samples = 1500
np.random.seed(0)
X, y= datasets.make_blobs(n_samples=2*n_samples, random_state=20,centers=num_classes)
# normalize dataset for easier parameter selection
X = StandardScaler().fit_transform(X)

X_train, X_test, y_train, y_test = train_test_split(
     X, y, test_size=0.33, random_state=42)



class MyDataset(Dataset):

    def __init__(self, X,y):
        self.x_data, self.y_data = X,y
    
    def __getitem__(self, i):
        return self.x_data[i], self.y_data[i]

    def __len__(self):
        return len(self.x_data)

dataset_train = MyDataset(X_train,y_train)
train_loader = DataLoader(dataset_train, batch_size=batch_size, shuffle=True)
dataset_test = MyDataset(X_test,y_test)
test_loader = DataLoader(dataset_test, batch_size=batch_size, shuffle=False)


class Classif():
    def __init__(self, clf):
        
        step_size = 1e-3
        self.num_features = 2
        self.num_classes  = 4
        self.batch_size   = 5

        self.init_nn, self.apply_nn = clf
        self.opt_init, self.opt_update, self.get_params = optimizers.adam(step_size)

    def loss(self, params, data, targets):
        preds = self.apply_nn(params, data)
        return -jnp.sum(preds * targets)

    @jit
    def update(self, params, x, y, opt_state):
        """ Compute the gradient for a batch and update the parameters """
        value, grads = value_and_grad(self.loss)(params, x, y)
        opt_state = self.opt_update(0, grads, opt_state)
        return self.get_params(opt_state), opt_state, value
    
    def accuracy(self, params, data_loader):
        acc_total = 0
        for batch_idx, (data, target) in enumerate(data_loader):
            data = jnp.array(data)
            targets = one_hot(np.array(target), self.num_classes)

            target_class = jnp.argmax(targets, axis=1)
            predicted_class = jnp.argmax(self.apply_nn(params, data), axis=1)
            acc_total += jnp.sum(predicted_class == target_class)
        return acc_total/len(data_loader.dataset)



    def run_training_loop(self, num_epochs, train_loader, test_loader, seed=42):
        """ Implements a learning loop over epochs. """
        # Initialize placeholder for loggin
        log_acc_train, log_acc_test, train_loss = [], [], []
        
        key = jax.random.PRNGKey(seed)
        params_init = self.init_nn(key, (self.batch_size, self.num_features))
        opt_state = self.opt_init(params_init)

        # Get the initial set of parameters
        params = self.get_params(opt_state)
    
        # Get initial accuracy after random init
        train_acc = self.accuracy(params, train_loader)
        test_acc = self.accuracy(params, test_loader)
        log_acc_train.append(train_acc)
        log_acc_test.append(test_acc)
    
        print(f"begin: train_acc {train_acc:0.4f}, test_acc {test_acc:0.4f} ")

        # Loop over the training epochs
        for epoch in range(num_epochs):
            start_time = time.time()
            for batch_idx, (data, target) in enumerate(train_loader):
                x = jnp.array(data)
                y = one_hot(np.array(target), self.num_classes)
                params, opt_state, loss = self.update(params, x, y, opt_state)
                train_loss.append(loss)

            epoch_time = time.time() - start_time
            train_acc = self.accuracy(params, train_loader)
            test_acc = self.accuracy(params, test_loader)
            log_acc_train.append(train_acc)
            log_acc_test.append(test_acc)
            print(f"Epoch {epoch+1} | T: {epoch_time:0.2f} | Train Acc: {train_acc:0.4f} | Test Acc: {test_acc:0.4f}")

        return train_loss, log_acc_train, log_acc_test, params, params_init

###### 
model = stax.serial(Dense(num_classes),
                        LogSoftmax
                      )
clf = Classif(model)

num_epochs=40
train_loss, log_acc_train, log_acc_test, params, params_init = \
        clf.run_training_loop(num_epochs, train_loader, test_loader)

It yields

TypeError: Argument '<__main__.Classif object at 0x7f5c990a47c0>' of type <class '__main__.Classif'> is not a valid JAX type.

if I switch off the @jit decorator then I get an error according to "interger" in value_and_grad, but my X, y and params are floats...

Now, everthing is ok is I don not use a python class at all... Is someone can help?
Thanks

[hawkinsp]

The "integer value" error from grad means what it says on the tin. If I print your params, I get:

((5, 4), [(DeviceArray([[ 0.02949349,  0.09376334, -0.79533563,  0.96614316],
             [ 0.59040182, -1.07658416, -0.10702543,  1.09999998]],            dtype=float64), DeviceArray([0.00272605, 0.01357591, 0.00355436, 0.00465731], dtype=float64)), ()])

Note the first two values are integers. Either don't pass them as part of params, change them to inexact types, or tell grad not to differentiate with respect to them.

The issue about "not a valid JAX value" is because of the self parameter.
If you want to jit a class method, try using something like:

@functools.partial(jit, static_argnums=0)

This requires that the class is immutable and has a well-defined __hash__ and __eq__, which does look to be the case here.

Hope that helps!

[jecampagne]

Thanks for your interest.
Well I'm a bit perplex as if I do not use Python class I get for the params for 5 blobs (sorry)

[(DeviceArray([[-0.28115242,  0.61787717, -0.46786687, -0.46523874,
               0.06839928],
             [-0.53766659, -0.31535681, -0.85428573, -0.08357738,
              -0.10971196]], dtype=float64), DeviceArray([-0.00453924, -0.0041792 ,  0.00483787,  0.00664846,
              0.01265203], dtype=float64)), ()]
begin: train_acc 0.22786069651741295, test_acc 0.21313131313131314 
Epoch 1 | T: 0.80 | Train Acc: 0.5990 | Test Acc: 0.5848

[jecampagne]

But effectively in the Python class the initilal params gets a couple of (5,5) ???

init params ((5, 5), [(DeviceArray([[-0.28115242, 0.61787717, -0.46786687, -0.46523874,
0.06839928],
[-0.53766659, -0.31535681, -0.85428573, -0.08357738,
-0.10971196]], dtype=float64), DeviceArray([-0.00453924, -0.0041792 , 0.00483787, 0.00664846,
0.01265203], dtype=float64)), ()])

[jecampagne]

In fact I got it for the integer stuff

   _,params_init = self.init_nn(key, (self.batch_size, self.num_features))

is simple solution.

And thanks to one of your remark I finnaly a code working :). The self argument makes the things tricky
So, when I use jit I must tell to skip the first arg...

 @functools.partial(jit, static_argnums=0)
    def update(self, params, x, y, opt_state):

Thanks