JAX port for Detectron2

[rozeappletree]

Hi, I have been using Detectron2 for almost a year for training large scale models and their results (mask rcnn) always fascinate me. Like FastAI it is one of the most flexible frameworks I have used for deeplearning.

I wanted to understand internals of both JAX and Detectron2 for some time and I think building flax port for Detectron2 will get me there.

Can I know if I can get any support from the team for solving issues if I get stuck.

Timeline:

6mos to 1year starting from Feb 2022. I have couple of commitments that hinder me from full time contribution but can spare weekends and some weekdays for the project. The idea is to slowly but effectively build the project.

Motivation:

Building mature frameworks for training models (stress on large-scale) is always critical for both research and industry. And JAX is currently the hottest topic I want to learn by doing.

[rozeappletree]

Linking issue from flax google/flax#1740