- Notifications
You must be signed in to change notification settings - Fork0
PyTorch image models, scripts, pretrained weights -- ResNet, ResNeXT, EfficientNet, EfficientNetV2, NFNet, Vision Transformer, MixNet, MobileNet-V3/V2, RegNet, DPN, CSPNet, and more
License
drinkingcoder/pytorch-image-models
Folders and files
Name | Name | Last commit message | Last commit date | |
---|---|---|---|---|
Repository files navigation
Universal feature extraction, new models, new weights, new test sets.
- All models support the
features_only=True
argument forcreate_model
call to return a network that extracts feature maps from the deepest layer at each stride. - New models
- CSPResNet, CSPResNeXt, CSPDarkNet, DarkNet
- ReXNet
- (Modified Aligned) Xception41/65/71 (a proper port of TF models)
- New trained weights
- SEResNet50 - 80.3
- CSPDarkNet53 - 80.1 top-1
- CSPResNeXt50 - 80.0 to-1
- DPN68b - 79.2 top-1
- EfficientNet-Lite0 (non-TF ver) - 75.5 (trained by @hal-314)
- Add 'real' labels for ImageNet and ImageNet-Renditions test set, see
results/README.md
- Train script and loader/transform tweaks to punch through more aug arguments
- README and documentation overhaul. See initial (WIP) documentation athttps://rwightman.github.io/pytorch-image-models/
Bunch of changes:
- DenseNet models updated with memory efficient addition from torchvision (fixed a bug), blur pooling and deep stem additions
- VoVNet V1 and V2 models added, 39 V2 variant (ese_vovnet_39b) trained to 79.3 top-1
- Activation factory added along with new activations:
- select act at model creation time for more flexibility in using activations compatible with scripting or tracing (ONNX export)
- hard_mish (experimental) added with memory-efficient grad, along with ME hard_swish
- context mgr for setting exportable/scriptable/no_jit states
- Norm + Activation combo layers added with initial trial support in DenseNet and VoVNet along with impl of EvoNorm and InplaceAbn wrapper that fit the interface
- Torchscript works for all but two of the model types as long as using Pytorch 1.5+, tests added for this
- Some import cleanup and classifier reset changes, all models will have classifier reset to nn.Identity on reset_classifer(0) call
- Prep for 0.1.28 pip release
- Add ResNeSt models (code adapted fromhttps://github.com/zhanghang1989/ResNeSt, paperhttps://arxiv.org/abs/2004.08955))
- Pruned EfficientNet B1, B2, and B3 (https://arxiv.org/abs/2002.08258) contributed byYonathan Aflalo
- Merged a number of execellent contributions in the ResNet model family over the past month
- BlurPool2D and resnetblur models initiated byChris Ha, I trained resnetblur50 to 79.3.
- TResNet models and SpaceToDepth, AntiAliasDownsampleLayer layers bymrT23
- ecaresnet (50d, 101d, light) models and two pruned variants using pruning as per (https://arxiv.org/abs/2002.08258) byYonathan Aflalo
- 200 pretrained models in total now with updated results csv in results folder
- Add some newly trained MobileNet-V2 models trained with latest h-params, rand augment. They compare quite favourably to EfficientNet-Lite
- 3.5M param MobileNet-V2 100 @ 73%
- 4.5M param MobileNet-V2 110d @ 75%
- 6.1M param MobileNet-V2 140 @ 76.5%
- 5.8M param MobileNet-V2 120d @ 77.3%
- Add EfficientNet-Lite models w/ weights ported fromTensorflow TPU
- Add RandAugment trained ResNeXt-50 32x4d weights with 79.8 top-1. Trained byAndrew Lavin (see Training section for hparams)
PyTorchImageModels (timm
) is a collection of image models, layers, utilities, optimizers, schedulers, data-loaders / augmentations, and reference training / validation scripts that aim to pull together a wide variety of SOTA models with ability to reproduce ImageNet training results.
The work of many others is present here. I've tried to make sure all source material is acknowledged via links to github, arxiv papers, etc in the README, documentation, and code docstrings. Please let me know if I missed anything.
All model architecture families include variants with pretrained weights. The are variants without any weights. Help training new or better weights is always appreciated. Here are some exampletraining hparams to get you started.
A full version of the list below with source links can be found in thedocumentation.
- CspNet (Cross-Stage Partial Networks) -https://arxiv.org/abs/1911.11929
- DenseNet -https://arxiv.org/abs/1608.06993
- DLA -https://arxiv.org/abs/1707.06484
- DPN (Dual-Path Network) -https://arxiv.org/abs/1707.01629
- EfficientNet (MBConvNet Family)
- EfficientNet NoisyStudent (B0-B7, L2) -https://arxiv.org/abs/1911.04252
- EfficientNet AdvProp (B0-B8) -https://arxiv.org/abs/1911.09665
- EfficientNet (B0-B7) -https://arxiv.org/abs/1905.11946
- EfficientNet-EdgeTPU (S, M, L) -https://ai.googleblog.com/2019/08/efficientnet-edgetpu-creating.html
- FBNet-C -https://arxiv.org/abs/1812.03443
- MixNet -https://arxiv.org/abs/1907.09595
- MNASNet B1, A1 (Squeeze-Excite), and Small -https://arxiv.org/abs/1807.11626
- MobileNet-V2 -https://arxiv.org/abs/1801.04381
- Single-Path NAS -https://arxiv.org/abs/1904.02877
- HRNet -https://arxiv.org/abs/1908.07919
- Inception-V3 -https://arxiv.org/abs/1512.00567
- Inception-ResNet-V2 and Inception-V4 -https://arxiv.org/abs/1602.07261
- MobileNet-V3 (MBConvNet w/ Efficient Head) -https://arxiv.org/abs/1905.02244
- NASNet-A -https://arxiv.org/abs/1707.07012
- PNasNet -https://arxiv.org/abs/1712.00559
- RegNet -https://arxiv.org/abs/2003.13678
- ResNet/ResNeXt
- ResNet (v1b/v1.5) -https://arxiv.org/abs/1512.03385
- ResNeXt -https://arxiv.org/abs/1611.05431
- 'Bag of Tricks' / Gluon C, D, E, S variations -https://arxiv.org/abs/1812.01187
- Weakly-supervised (WSL) Instagram pretrained / ImageNet tuned ResNeXt101 -https://arxiv.org/abs/1805.00932
- Semi-supervised (SSL) / Semi-weakly Supervised (SWSL) ResNet/ResNeXts -https://arxiv.org/abs/1905.00546
- ECA-Net (ECAResNet) -https://arxiv.org/abs/1910.03151v4
- Squeeze-and-Excitation Networks (SEResNet) -https://arxiv.org/abs/1709.01507
- Res2Net -https://arxiv.org/abs/1904.01169
- ResNeSt -https://arxiv.org/abs/2004.08955
- ReXNet -https://arxiv.org/abs/2007.00992
- SelecSLS -https://arxiv.org/abs/1907.00837
- Selective Kernel Networks -https://arxiv.org/abs/1903.06586
- TResNet -https://arxiv.org/abs/2003.13630
- VovNet V2 and V1 -https://arxiv.org/abs/1911.06667
- Xception -https://arxiv.org/abs/1610.02357
- Xception (Modified Aligned, Gluon) -https://arxiv.org/abs/1802.02611
- Xception (Modified Aligned, TF) -https://arxiv.org/abs/1802.02611
Several (less common) features that I often utilize in my projects are included. Many of their additions are the reason why I maintain my own set of models, instead of using others' via PIP:
- All models have a common default configuration interface and API for
- accessing/changing the classifier -
get_classifier
andreset_classifier
- doing a forward pass on just the features -
forward_features
(seedocumentation) - these makes it easy to write consistent network wrappers that work with any of the models
- accessing/changing the classifier -
- All models support multi-scale feature map extraction (feature pyramids) via create_model (seedocumentation)
create_model(name, features_only=True, out_indices=..., output_stride=...)
out_indices
creation arg specifies which feature maps to return, these indices are 0 based and generally correspond to theC(i + 1)
feature level.output_stride
creation arg controls output stride of the network by using dilated convolutions. Most networks are stride 32 by default. Not all networks support this.- feature map channel counts, reduction level (stride) can be queried AFTER model creation via the
.feature_info
member
- All models have a consistent pretrained weight loader that adapts last linear if necessary, and from 3 to 1 channel input if desired
- High performancereference training, validation, and inference scripts that work in several process/GPU modes:
- NVIDIA DDP w/ a single GPU per process, multiple processes with APEX present (AMP mixed-precision optional)
- PyTorch DistributedDataParallel w/ multi-gpu, single process (AMP disabled as it crashes when enabled)
- PyTorch w/ single GPU single process (AMP optional)
- A dynamic global pool implementation that allows selecting from average pooling, max pooling, average + max, or concat([average, max]) at model creation. All global pooling is adaptive average by default and compatible with pretrained weights.
- A 'Test Time Pool' wrapper that can wrap any of the included models and usually provide improved performance doing inference with input images larger than the training size. Idea adapted from original DPN implementation when I ported (https://github.com/cypw/DPNs)
- Learning rate schedulers
- Ideas adopted from
- AllenNLP schedulers
- FAIRseq lr_scheduler
- SGDR: Stochastic Gradient Descent with Warm Restarts (https://arxiv.org/abs/1608.03983)
- Schedulers include
step
,cosine
w/ restarts,tanh
w/ restarts,plateau
- Ideas adopted from
- Optimizers:
rmsprop_tf
adapted from PyTorch RMSProp by myself. Reproduces much improved Tensorflow RMSProp behaviour.radam
byLiyuan Liu (https://arxiv.org/abs/1908.03265)novograd
byMasashi Kimura (https://arxiv.org/abs/1905.11286)lookahead
adapted from impl byLiam (https://arxiv.org/abs/1907.08610)fused<name>
optimizers by name withNVIDIA Apex installedadamp
andsgdp
byNaver ClovAI (https://arxiv.org/abs/2006.08217)
- Random Erasing fromZhun Zhong (https://arxiv.org/abs/1708.04896)
- Mixup (as inhttps://arxiv.org/abs/1710.09412) - currently implementing/testing
- An inference script that dumps output to CSV is provided as an example
- AutoAugment (https://arxiv.org/abs/1805.09501) and RandAugment (https://arxiv.org/abs/1909.13719) ImageNet configurations modeled after impl for EfficientNet training (https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/autoaugment.py)
- AugMix w/ JSD loss (https://arxiv.org/abs/1912.02781), JSD w/ clean + augmented mixing support works with AutoAugment and RandAugment as well
- SplitBachNorm - allows splitting batch norm layers between clean and augmented (auxiliary batch norm) data
- DropPath aka "Stochastic Depth" (https://arxiv.org/abs/1603.09382)
- DropBlock (https://arxiv.org/abs/1810.12890)
- Efficient Channel Attention - ECA (https://arxiv.org/abs/1910.03151)
- Blur Pooling (https://arxiv.org/abs/1904.11486)
- Space-to-Depth bymrT23 (https://arxiv.org/abs/1801.04590) -- original paper?
Model validation results can be found in thedocumentation and in theresults tables
See thedocumentation
About
PyTorch image models, scripts, pretrained weights -- ResNet, ResNeXT, EfficientNet, EfficientNetV2, NFNet, Vision Transformer, MixNet, MobileNet-V3/V2, RegNet, DPN, CSPNet, and more
Resources
License
Uh oh!
There was an error while loading.Please reload this page.
Stars
Watchers
Forks
Packages0
Languages
- Python99.9%
- Shell0.1%