Flax model much slower than PyTorch when Inferencing (Not utilized GPU fully)

[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)

[cgarciae]

Hey @Tharinda-EV, its a very nice benchmark!

1. Try profiling the JAX program, maybe Flax is using some inefficient operations and we can try to improve our layers.
2. Can you please also post this issue in the JAX repo? Flax ultimately is just JAX, and most of the performance gains/losses ultimately comes from the code that XLA produces (which we have very little control of sadly). If you manage to profile the program make sure to post it as well.