Package: jimm
JAX Image ModelsFuture home of JAX Image Models (jimm). Sibling library ofhttps://github.com/rwightman/pytorch-image-models. Liketimm,jimmwill become a collection of JAX based image models w/ pretrained weights focused on transfer learning. The first models to be include here will the Flax Linen JAX adaptation of the MBConv family (EfficientNet, MobileNetV2/V3, etc)https://github.com/rwightman/efficientnet-jax.jimmwill be built while exploring transfer learning and the impact of different augmentation, regularization, supervised, semi-supervised, self-supervised pretraining techniques on the transferability of weights for different target tasks.Specifically I hope to compare transfer learning on a wide variety of models for a variety of target datasets and tasks across:supervised ImageNet-1k training with heavy augmentation and regularizationsupervised ImageNet-1k training with light augmentation and regularizationsupervised 'larger' dataset (ImageNet-21k, OpenImages) training w/ light (may try heavy since this isn't JFT-300M) augmentation and regularizationsemi-supervised and self-supervised pretraining (SimCLRV2, BYOL, FixMatch, etc) on ImageNet-1k/21k, OpenImagesThere is already a body of research work on the subject of transfer learning. Much of my work here will not be breaking new ground but providing me with an opportunity to learn and do what I do best -- refine and improve. Some papers in this space:A Large-scale Study of Representation Learning with the Visual Task Adaptation Benchmark -https://arxiv.org/abs/1910.04867Big Transfer (BiT): General Visual Representation Learning-https://arxiv.org/abs/1912.11370On Robustness and Transferability of Convolutional Neural Networks -https://arxiv.org/abs/2007.08558Which Model to Transfer? Finding the Needle in the Growing Haystack -https://arxiv.org/abs/2010.06402Self-supervised Pre-training with Hard Examples Improves Visual Representations -https://arxiv.org/abs/2012.13493Do Adversarially Robust ImageNet Models Transfer Better? -https://arxiv.org/abs/2007.08489How Useful is Self-Supervised Pretraining for Visual Tasks? -https://arxiv.org/abs/2003.14323... please file an issue if you have ideas for additional papersPapers on self or semi-supervised techniques that I plan to exploreBootstrap your own latent: A new approach to self-supervised Learning -https://arxiv.org/abs/2006.07733Big Self-Supervised Models are Strong Semi-Supervised Learners -https://arxiv.org/abs/2006.10029FixMatch: Simplifying Semi-Supervised Learning with Consistency and Confidence -https://arxiv.org/abs/2001.07685Self-training with Noisy Student improves ImageNet classification -https://arxiv.org/abs/1911.04252... please file an issue if you have ideas for additional papersThe scope of 'transfer learning' will initially cover fine-tuning (head replaced with newly initialized task-specific head, 0..N layers frozen). I may explore linear classifier or low-shot techniques from semi/self-supervised pretraining later.I'm currently planning which datasets to select for transfer learning benchmarks. I'm hoping to explore a cross section of natural image datasets (that don't overlap with imagenet), and other (medical, spectrogram, industrial inspection, etc). I'd like to cover a different cross section of datasets than usual, I hope to find some interesting options from various Kaggle challenges, etc that are available with compatible licenses. Dataset suggestions welcome.The development ofjimmdoes not mean I'm abandoningtimm. I will build the models in a manner that allows easy movements of weights back and forth. I'm interested in building in JAX because a) I've enjoyed my JAX exploration so far b) I have some TPU credits that allows more compute intensive exploration than my open source training budget would allow. I will be augmenting that with my local NVIDIA GPU resources.