PyTorch image models, scripts, pretrained weights -- ResNet, ResNeXT, EfficientNet, EfficientNetV2, NFNet, Vision Transformer, MixNet, MobileNet-V3/V2, RegNet, DPN, CSPNet, and more
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
Go to file
Ross Wightman dfe80414a6
Add bool arg helper
4 years ago
.github/workflows autosquash github workflow didn't work out, removing 4 years ago
convert Add working Gluon Xception-65 model. Some cleanup still needed. 5 years ago
docs Continuing README and documentation updates 4 years ago
notebooks ImageNet-1k vs ImageNet-v2 comparison 6 years ago
results Re-ran batch validation on all models across all sets 4 years ago
tests More uniform treatment of classifiers across all models, reduce code duplication. 4 years ago
timm Add bool arg helper 4 years ago
.gitattributes Add .gitattributes 6 years ago
.gitignore Implement Eca modules 5 years ago
LICENSE Add Apache LICENSE file 6 years ago
MANIFEST.in Fix pruned txt files not being installed during pip install 5 years ago
README.md Continuing README and documentation updates 4 years ago
avg_checkpoints.py Fix some attributions, add copyrights to some file docstrings 4 years ago
clean_checkpoint.py Fix some attributions, add copyrights to some file docstrings 4 years ago
distributed_train.sh Fix distributed train script 6 years ago
hubconf.py Simpler approach to loading entrypoints in hubconf works properly 5 years ago
inference.py Merge branch 'logger' into features. Change 'logger' to '_logger'. 4 years ago
mkdocs.yml Continuing README and documentation updates 4 years ago
requirements-docs.txt configure mkdocs based documentation 5 years ago
requirements-sotabench.txt Update requirements so PyTorch 1.4 is min, add separate sotabench reqs to force 1.4 and prevent error with 1.5 5 years ago
requirements.txt Update requirements so PyTorch 1.4 is min, add separate sotabench reqs to force 1.4 and prevent error with 1.5 5 years ago
setup.py Update setup.py 5 years ago
sotabench.py Update README with model results and attribution. Make scheduler factory bit more robust to arg differences, add noise to plateau lr and fix min/max. 5 years ago
sotabench_setup.sh Update requirements so PyTorch 1.4 is min, add separate sotabench reqs to force 1.4 and prevent error with 1.5 5 years ago
train.py Add more augmentation arguments, including a no_aug disable flag. Fix #209 4 years ago
validate.py Update results csv file rank/diff script and small validate script tweak for batch validation 4 years ago

README.md

PyTorch Image Models, etc

What's New

Aug 1, 2020

Universal feature extraction, new models, new weights, new test sets.

  • All models support the features_only=True argument for create_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

June 11, 2020

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

May 12, 2020

May 3, 2020

May 1, 2020

  • Merged a number of execellent contributions in the ResNet model family over the past month
  • 200 pretrained models in total now with updated results csv in results folder

April 5, 2020

  • 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%

March 18, 2020

  • Add EfficientNet-Lite models w/ weights ported from Tensorflow TPU
  • Add RandAugment trained ResNeXt-50 32x4d weights with 79.8 top-1. Trained by Andrew Lavin (see Training section for hparams)

Introduction

PyTorch Image Models (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.

Models

Most included models have pretrained weights. The weights are either from their original sources, ported by myself from their original framework (e.g. Tensorflow models), or trained from scratch using the included training script. A full version of the list below with source links and references can be found in the documentation.

Features

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 and reset_classifier
    • doing a forward pass on just the features - forward_features
    • these makes it easy to write consistent network wrappers that work with any of the models
  • All models support multi-scale feature map extraction (feature pyramids) via create_model
    • 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 the C(i + 1) feature level. Most models start with stride 2 features (C1) at index 0 and end with C5 at index 4. Some models start with stride 1 or 4 and end with 6 (stride 64).
    • output_stride creation arg controls output stride of the network, most networks are stride 32 by default. Dilated convs are used to limit the output stride. 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 performance reference 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
  • Optimizers:
  • Random Erasing from Zhun Zhong (https://arxiv.org/abs/1708.04896)
  • Mixup (as in https://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 by mrT23 (https://arxiv.org/abs/1801.04590) -- original paper?

Results

Model validation results can be found in the documentation and in the results tables

Getting Started

See the documentation