EfficientTrain icon indicating copy to clipboard operation
EfficientTrain copied to clipboard

1.5−3.0× lossless training or pre-training speedup. An off-the-shelf, easy-to-implement algorithm for the efficient training of foundation visual backbones.

EfficientTrain (ICCV 2023)

This repo is used to release the code and pre-trained models for the EfficientTrain algorithm.

EfficientTrain: Exploring Generalized Curriculum Learning for Training Visual Backbones. ICCV 2023.
Yulin Wang*, Yang Yue*, Rui Lu, Tianjiao Liu, Zhao Zhong, Shiji Song and Gao Huang
Tsinghua University, Huawei, BAAI
[arXiv]

Overview

In this paper, we present a novel curriculum learning approach for the efficient training of visual backbones. Our algorithm, EfficientTrain, is simple, general, yet surprisingly effective. For example, it reduces the training time of a wide variety of popular models (e.g., ConvNeXts, DeiT, PVT, and Swin/CSWin Transformers) by more than 1.5× on ImageNet-1K/22K without sacrificing the accuracy. It is effective for self-supervised learning (i.e., MAE) as well.

Catalog

  • [x] ImageNet-1K Training Code
  • [x] ImageNet-1K Pre-trained Models

Installation

We support Pytorch>=1.8.0 and torchvision>=0.9.0. Please install them following official instructions. An example is given in the following.

pip install torch==1.8.0+cu111 torchvision==0.9.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html

Clone this repo and install required packages:

git clone https://github.com/LeapLabTHU/EfficientTrain
pip install timm==0.3.2 tensorboardX six

The instructions for preparing ImageNet-1K/22K datasets can be found here.

Training

See TRAINING.md for the training instructions.

Pre-trained Models & Evaluation

See EVAL.md for the pre-trained models and the evaluation instructions.

Results

Supervised learning on ImageNet-1K

ImageNet-22K pre-training

Supervised learning on ImageNet-1K (varying training epochs and test input sizes)

Object detection and instance segmentation on COCO

Self-supervised learning results on top of MAE

TODO

This repo is still being updated.

  • [ ] ImageNet-22K Training Code
  • [ ] ImageNet-22K Pre-trained Models
  • [ ] ImageNet-1K Self-supervised Learning Code (EfficientTrain + MAE)
  • [ ] EfficientTrain + MAE Pre-trained Models

Acknowledgement

This repo is mainly developed on the top of ConvNeXt, we sincerely thank them for their efficient and neat codebase. This repo is also built using DeiT and timm.

Citation

If you find this repo helpful, please consider citing us.

@InProceedings{wang2023efficienttrain,
    title={EfficientTrain: Exploring Generalized Curriculum Learning for Training Visual Backbones},
   author={Wang, Yulin and Yue, Yang and Lu, Rui and Liu, Tianjiao and Zhong, Zhao and Song, Shiji and Huang, Gao},
 booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)},
     year={2023}
}