How to quantize the pre trained JAX model bsed on flax.linen

[deshwalmahesh]

I have a pre-trained JAX model for MAXIM: Image Enhancement. Now to reduce the runtime and use it in production, I'll have to quantize the weights. I have 2 options since there is no direct conversion to ONNX.

1. JAX -> Tensorflow -> ONNX (Help Thread)
2. JAX -> TFLite

Going for the second option, there's this function tf.lite.TFLiteConverter.experimental_from_jax

Looking at this official example, the code block

serving_func = functools.partial(predict, params)
x_input = jnp.zeros((1, 28, 28))
converter = tf.lite.TFLiteConverter.experimental_from_jax(
    [serving_func], [[('input1', x_input)]])
tflite_model = converter.convert()
with open('jax_mnist.tflite', 'wb') as f:
  f.write(tflite_model)

it seems to be using the params from the model and a function predict which in case are defined while model building and training itself as

predict:

init_random_params, predict = stax.serial(
    stax.Flatten,
    stax.Dense(1024), stax.Relu,
    stax.Dense(1024), stax.Relu,
    stax.Dense(10), stax.LogSoftmax)

and the params

opt_init, opt_update, get_params = optimizers.momentum(step_size, mass=momentum_mass)
params = get_params(opt_state)

My question is that how can I get these two required params and predict for my pre-trained model so that I can try to replicate example for my own model?

[marcvanzee]

Hi @deshwalmahesh, jax2tf provides an example of converting a Flax model to TFLite, which seems to be closer to your use-case: https://github.com/google/jax/tree/main/jax/experimental/jax2tf/examples/tflite/mnist. I would recommend converting using this path, since we provide more support than the tf.lite.experimental_from_jax path.

This path first converts the JAX graph to a TF graph using jax2tf, and then converts the TF graph to a TFLite model using the TFLite converter. We are actually planning on adding this path to tf.lite.experimental_from_jax under a flag soon, but this may take a few weeks. Conceptually both paths are very similar, the jax2tf path currently just take a bit more code, but it has better support.

If you encounter any non-Flax related problems with this conversion path please file an issue against the JAX repo starting with "[jax2tf] ....", we will try resolving your issue there then.

[Ir1d]

Hi @deshwalmahesh , did you succeed in converting MAXIM to ONNX? Could you please share your conversion script or converted model?

[deshwalmahesh]

Hi! I actually found a way and I'm not sure if it's the right one but it's working for me.

So I got an answer on the official repo. Here is the code:

import tensorflow as tf
from jax.experimental import jax2tf


def predict(input_img):
  '''
  Function to predict the output from the JAX model
  '''
  return model.apply({'params': flax.core.freeze(params)}, input_img)


tf_predict = tf.function(
    jax2tf.convert(predict, enable_xla=False),
    input_signature=[
        tf.TensorSpec(shape=[1, 704, 1024, 3], dtype=tf.float32, name='input_image')
    ],
    autograph=False)

converter = tf.lite.TFLiteConverter.from_concrete_functions(
    [tf_predict.get_concrete_function()], tf_predict)

converter.target_spec.supported_ops = [
    tf.lite.OpsSet.TFLITE_BUILTINS,  # enable TensorFlow Lite ops.
    tf.lite.OpsSet.SELECT_TF_OPS  # enable TensorFlow ops.
]
tflite_float_model = converter.convert()

with open('float_model.tflite', "wb") as f: f.write(tflite_float_model)


converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_quantized_model = converter.convert()

with open('./quantized.tflite', 'wb') as f: f.write(tflite_quantized_model)

You can now load and run the model easily using tf.lite.Interpreter