Merge pull request #323 from rwightman/imagenet21k_datasets_more

BiT (Big Transfer) ResNetV2 models, Official ViT Hybrid R50 weights, VIT IN21K weights updated w/ repr layer, ImageNet21k and dataset / parser refactor

Author: rwightman

README.md

@@ -2,6 +2,19 @@
 
 ## What's New
 
+### Jan 25, 2021
+* Add ResNetV2 Big Transfer (BiT) models w/ ImageNet-1k and 21k weights from https://github.com/google-research/big_transfer
+* Add official R50+ViT-B/16 hybrid models + weights from https://github.com/google-research/vision_transformer
+* ImageNet-21k ViT weights are added w/ model defs and representation layer (pre logits) support
+  * NOTE: ImageNet-21k classifier heads were zero'd in original weights, they are only useful for transfer learning
+* Add model defs and weights for DeiT Vision Transformer models from https://github.com/facebookresearch/deit
+* Refactor dataset classes into ImageDataset/IterableImageDataset + dataset specific parser classes
+* Add Tensorflow-Datasets (TFDS) wrapper to allow use of TFDS image classification sets with train script
+  * Ex: `train.py /data/tfds --dataset tfds/oxford_iiit_pet --val-split test --model resnet50 -b 256 --amp --num-classes 37 --opt adamw --lr 3e-4 --weight-decay .001 --pretrained -j 2`
+* Add improved .tar dataset parser that reads images from .tar, folder of .tar files, or .tar within .tar
+  * Run validation on full ImageNet-21k directly from tar w/ BiT model: `validate.py /data/fall11_whole.tar --model resnetv2_50x1_bitm_in21k --amp`
+* Models in this update should be stable w/ possible exception of ViT/BiT, possibility of some regressions with train/val scripts and dataset handling
+
 ### Jan 3, 2021
 * Add SE-ResNet-152D weights
   * 256x256 val, 0.94 crop top-1 - 83.75
@@ -130,7 +143,9 @@ All model architecture families include variants with pretrained weights. The ar
 
 A full version of the list below with source links can be found in the [documentation](https://rwightman.github.io/pytorch-image-models/models/).
 
+* Big Transfer ResNetV2 (BiT) - https://arxiv.org/abs/1912.11370
 * CspNet (Cross-Stage Partial Networks) - https://arxiv.org/abs/1911.11929
+* DeiT (Vision Transformer) - https://arxiv.org/abs/2012.12877
 * DenseNet - https://arxiv.org/abs/1608.06993
 * DLA - https://arxiv.org/abs/1707.06484
 * DPN (Dual-Path Network) - https://arxiv.org/abs/1707.01629
@@ -242,6 +257,10 @@ One of the greatest assets of PyTorch is the community and their contributions.
 * Albumentations - https://github.com/albumentations-team/albumentations
 * Kornia - https://github.com/kornia/kornia
 
+### Knowledge Distillation
+* RepDistiller - https://github.com/HobbitLong/RepDistiller
+* torchdistill - https://github.com/yoshitomo-matsubara/torchdistill
+
 ### Metric Learning
 * PyTorch Metric Learning - https://github.com/KevinMusgrave/pytorch-metric-learning
 

docs/models.md

@@ -10,6 +10,10 @@ Most included models have pretrained weights. The weights are either:
 
 The validation results for the pretrained weights can be found [here](results.md)
 
+## Big Transfer ResNetV2 (BiT) [[resnetv2.py](https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/resnetv2.py)]
+* Paper: `Big Transfer (BiT): General Visual Representation Learning` - https://arxiv.org/abs/1912.11370
+* Reference code: https://github.com/google-research/big_transfer
+
 ## Cross-Stage Partial Networks [[cspnet.py](https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/cspnet.py)]
 * Paper: `CSPNet: A New Backbone that can Enhance Learning Capability of CNN` - https://arxiv.org/abs/1911.11929
 * Reference impl: https://github.com/WongKinYiu/CrossStagePartialNetworks

inference.py

@@ -13,7 +13,7 @@
 import torch
 
 from timm.models import create_model, apply_test_time_pool
-from timm.data import Dataset, create_loader, resolve_data_config
+from timm.data import ImageDataset, create_loader, resolve_data_config
 from timm.utils import AverageMeter, setup_default_logging
 
 torch.backends.cudnn.benchmark = True
@@ -83,7 +83,7 @@ def main():
         model = model.cuda()
 
     loader = create_loader(
-        Dataset(args.data),
+        ImageDataset(args.data),
         input_size=config['input_size'],
         batch_size=args.batch_size,
         use_prefetcher=True,