Flax model much slower than PyTorch when Inferencing

[vidulae]

Hi,

I am running the same ResNet50 model with the same weights using Pytorch and Flax, but I am seeing quite a bad performance for Flax. And also when inferencing the PyTorch model, GPU utilization is around 100% but in the FLAX model, GPU utilization is about around 35%. I checked the arrays also, they are stored in GPU.

Model	Inferencing Time per loop (1 iteration)	Inferencing Time per loop (100 iterations)
Pytorch (GPU)	~12 ms	~11 ms
Flax (GPU)	~66 ms	~60 ms

I have warmed the jit function, and I think I have put the model and data on the device in both cases. What am I doing wrong here? Is there anything optimized way to use jit functions in building FLAX models?

Notebook: https://colab.research.google.com/drive/1b486yGovsLLuGawwhFmp6r6YUNw6Eupo?usp=sharing

Environment info

• Platform: Ubuntu-20.04
• Python version: 3.9.6
• PyTorch version (GPU?): 1.13.1+cu116 (True)
• Flax version (CPU?/GPU?/TPU?): 0.6.7 (GPU) (Installed using pip install flax)
• Jax version: 0.4.4
• JaxLib version: 0.4.4
• Using GPU in the script?: Yes
• Using distributed or parallel set-up in the script?: No
• Both Jax and Jaxlib are installed using: pip install --upgrade jax==0.4.4 jaxlib==0.4.4+cuda11.cudnn82 -f https://storage.googleapis.com/jax releases/jax_cuda_releases.html
• GPU: Nvidia GPU (Quadro RTX 6000)

[jakevdp]

cross-reference to your duplicate question in google/flax#2964, where I suspect you'll get a better answer.