EfficientNetV2 in Jax

[faresobeid]

I wanted to use a pre-trained efficientnetv2 model in flax but have found no easy way to do this. I'm using kaggle and want to import the pre-trained model and finetune it. What is the easiest way to do this?

[patrick-kidger]

I believe this is supported in Eqxvision: https://eqxvision.readthedocs.io/en/latest/api/models/classification/efficientnet_v2/

[faresobeid]

Thank you for the reply, I've just looked into equinox and am impressed by its features, I will try and rewrite my code in equinox and see if I prefer it over Flax.

[patrick-kidger]

Great! Heads-up that Eqxvision currently requires equinox==0.10.2, by the way. We're hoping to update it to the latest version shortly.

Also note that eqxvision uses both JAX and PyTorch (to get the pre-trained weights), and on some machines importing both JAX and PyTorch can cause issues, as they try to link against different versions of the same shared libraries. This can usually be solved by either doing:

import jax
jax.random.PRNGKey(0)  # initialise the JAX backend by doing a JAX op
import torch

before the rest of your script, or by doing:

import torch
import jax
jax.random.PRNGKey(0)

before the rest of your script. (Sometimes the import order matters.)

[ASEM000]

or keras with jax backend, I think you can install via pip install keras_core and do something like this import os; os.environ['KERAS_BACKEND']='jax'. soon new keras will support jax out of the box, and then you can install via pip install keras.

[faresobeid]

Thank you for replying. I will probably just move my code to equinox but if I decide to go back I will definitely give this a try!