SOTA Image Classification Models in PyTorch

Features

• Applicable for the following tasks:
  • Fine-tuning with custom classification datasets.
  • Used as a backbone in downstream tasks like object detection, semantic segmentation, pose estimation, etc.
• Almost no dependency in model usage.
• 10+ High-precision and High-efficient SOTA models.
• Regularly updated with new models.
• PyTorch, ONNX, CoreML, TFLite, OpenVINO Inference and Export.

Supported Models

2021

• ResT (NeurIPS 2021) (wofmanaf)
• MicroNet (ICCV 2021) (Yunsheng Li)
• ResNet* (ArXiv 2021) (Ross Wightman)
• PatchConvnet (ArXiv 2021) (Meta Research)

2022

• CSWin (CVPR 2022) (Microsoft)
• PVTv2 (CVMJ 2022) (whai362)
• UniFormer (ICLR 2022) (SenseTime X-Lab)
• PoolFormer (CVPR 2022) (Sea AI Lab)
• ConvNeXt (CVPR 2022) (Meta Research)
• VAN (ArXiv 2022) (BNRist)
• FocalNet (ArXiv 2022) (Microsoft)
• WaveMLP (CVPR 2022) (HUAWEI Noah's Ark Lab)
• DaViT (ArXiv 2022) (dingmyu)
• NAT (ArXiv 2022) (SHI Lab)
• FAN (ArXiv 2022) (NVlabs)
• ResTv2 (ArXiv 2022) (wofmanaf)

Losses

• PolyLoss (ICLR 2022) (yiyixuxu)

Benchmarks

Notes: ResNet* is from "ResNet strikes back" paper.

Usage

Requirements

• torch >= 1.11
• torchvision >= 0.12

Other requirements can be installed with pip install -r requirements.txt.

Show Supported Models

$ python list_models.py

A table with model names and variants will be shown:

                Supported Models
               
  Model Names  │ Model Variants
╶──────────────┼──────────────────────────────────╴
  ResNet       │ ['18', '34', '50', '101', '152']
  MicroNet     │ ['M1', 'M2', 'M3']
  ConvNeXt     │ ['T', 'S', 'B']
  VAN          │ ['S', 'B', 'L']
  PVTv2        │ ['B1', 'B2', 'B3', 'B4', 'B5']
  ResT         │ ['S', 'B', 'L']
  CSWin        │ ['T', 'S', 'B', 'L']
  WaveMLP      │ ['T', 'S', 'M']
  PoolFormer   │ ['S24', 'S36', 'M36']
  PatchConvnet │ ['S60', 'S120', 'B60']
  UniFormer    │ ['S', 'B']
  FocalNet     │ ['T', 'S', 'B']

Inference

# Example with VAN-S
$ python infer.py --source assests/dog.jpg --model VAN --variant S --checkpoint /path/to/van_s

You will see an output similar to this:

assests\dog.jpg >>>>> Golden retriever

Note: The above code is only for ImageNet pre-trained models. Modify the model's checkpoint loading and class names in infer.py for your custom needs.

Finetune

You can use any dataset from torchvision.datasets. For custom datasets, ImageFolder can be used to create a dataset class.

In this repo, finetuning on CIFAR-10 is provided in finetune.py.

!! What is not available yet:

• Distributed training
• Mixup and Cutmix augmentation

$ python finetune.py --cfg configs/finetune.yaml

Convert to other Frameworks

Install respective libraries for your desire framework:

# ONNX
$ pip install onnx onnx-simplifier onnxruntime
# CoreML
$ pip install coremltools
# OpenVINO
$ pip install onnx onnx-simplifier openvino-dev 
# TFLite (Coming Soon)
$ pip install onnx onnx-simplifier openvino-dev openvino2tensorflow tflite-runtime

Convert:

# ONNX
$ python convert/to_onnx.py --model MODEL_NAME --variant MODEL_VARIANT --num_classes NUM_CLASSES --checkpoint /path/to/weights --size IMAGE_SIZE

# CoreML
$ python convert/to_coreml.py --model MODEL_NAME --variant MODEL_VARIANT --num_classes NUM_CLASSES --checkpoint /path/to/weights --size IMAGE_SIZE

# OpenVINO
$ python convert/to_openvino.py --model MODEL_NAME --variant MODEL_VARIANT --num_classes NUM_CLASSES --checkpoint /path/to/weights --size IMAGE_SIZE --precision FP32 or FP16

Inference:

# PyTorch
$ python convert/infer_pt.py --source IMG_FILE_PATH --model MODEL_NAME --variant MODEL_VARIANT --num_classes NUM_CLASSES --checkpoint /path/to/weights --size IMAGE_SIZE --device cuda or cpu

# ONNX
$ python convert/infer_onnx.py --source IMG_FILE_PATH --model MODEL_PATH

# OpenVINO
$ python convert/infer_openvino.py --source IMG_FILE_PATH --model MODEL_PATH --device CPU or GPU

Framework Comparison

CPU:

GPU:

Latency in milliseconds. Tested with Ryzen 7 4800HS and RTX 1650ti.

Acknowledgements

Most of the codes are borrowed from timm and DeiT. I would like to thank the papers' authors for open-sourcing their codes and providing pre-trained models.