This repository contains the official code and pretrained models for CoaT: Co-Scale Conv-Attentional Image Transformers. It introduces (1) a co-scale mechanism to realize fine-to-coarse, coarse-to-fine and cross-scale attention modeling and (2) an efficient conv-attention module to realize relative position encoding in the factorized attention.
For more details, please refer to CoaT: Co-Scale Conv-Attentional Image Transformers by Weijian Xu*, Yifan Xu*, Tyler Chang, and Zhuowen Tu.
-
Classification (ImageNet dataset)
Name Acc@1 Acc@5 #Params CoaT-Lite Tiny 77.5 93.8 5.7M CoaT-Lite Mini 79.1 94.5 11M CoaT-Lite Small 81.9 95.5 20M CoaT-Lite Medium 83.6 96.7 45M CoaT Tiny 78.3 94.0 5.5M CoaT Mini 81.0 95.2 10M CoaT Small 82.1 96.1 22M -
Instance Segmentation (Mask R-CNN w/ FPN on COCO dataset)
Name Schedule Bbox AP Segm AP CoaT-Lite Mini 1x 41.4 38.0 CoaT-Lite Mini 3x 42.9 38.9 CoaT-Lite Small 1x 45.2 40.7 CoaT-Lite Small 3x 45.7 41.1 CoaT Mini 1x 45.1 40.6 CoaT Mini 3x 46.5 41.8 CoaT Small 1x 46.5 41.8 CoaT Small 3x 49.0 43.7 -
Object Detection (Deformable-DETR on COCO dataset)
Name AP AP50 AP75 APS APM APL CoaT-Lite Small 47.0 66.5 51.2 28.8 50.3 63.3
09/27/2021: Code and pre-trained checkpoints for instance segmentation with MMDetection are released.
08/27/2021: Pre-trained checkpoints for CoaT Small and CoaT-Lite Medium are released.
05/19/2021: Pre-trained checkpoints for Mask R-CNN benchmark with CoaT-Lite Small backbone are released.
05/19/2021: Code and pre-trained checkpoints for Deformable-DETR with for CoaT-Lite Small backbone are released.
05/11/2021: Pre-trained checkpoints for CoaT-Lite Small are released.
05/09/2021: Pre-trained checkpoints for Mask R-CNN benchmark with CoaT Mini backbone are released.
05/06/2021: Pre-trained checkpoints for CoaT Mini are released.
05/02/2021: Pre-trained checkpoints for CoaT Tiny are released.
04/25/2021: Code and pre-trained checkpoints for Mask R-CNN benchmark with CoaT-Lite Mini backbone are released.
04/23/2021: Pre-trained checkpoints for CoaT-Lite Mini are released.
04/22/2021: Code and pre-trained checkpoints for CoaT-Lite Tiny are released.
The following usage is provided for the classification task using CoaT model. For the other tasks, please follow the corresponding readme, such as instance segmentation and object detection.
-
Set up a new conda environment and activate it.
# Create an environment with Python 3.8. conda create -n coat python==3.8 conda activate coat
-
Install required packages.
# Install PyTorch 1.7.1 w/ CUDA 11.0. pip install torch==1.7.1+cu110 torchvision==0.8.2+cu110 torchaudio==0.7.2 -f https://download.pytorch.org/whl/torch_stable.html # Install timm 0.3.2. pip install timm==0.3.2 # Install einops. pip install einops
-
Clone the repo.
git clone https://github.com/mlpc-ucsd/CoaT cd CoaT
-
Download ImageNet dataset (ILSVRC 2012) and extract.
# Create dataset folder. mkdir -p ./data/ImageNet # Download the dataset (not shown here) and copy the files (assume the download path is in $DATASET_PATH). cp $DATASET_PATH/ILSVRC2012_img_train.tar $DATASET_PATH/ILSVRC2012_img_val.tar $DATASET_PATH/ILSVRC2012_devkit_t12.tar.gz ./data/ImageNet # Extract the dataset. python -c "from torchvision.datasets import ImageNet; ImageNet('./data/ImageNet', split='train')" python -c "from torchvision.datasets import ImageNet; ImageNet('./data/ImageNet', split='val')" # After the extraction, you should observe `train` and `val` folders under ./data/ImageNet.
We provide the CoaT checkpoints pre-trained on the ImageNet dataset.
Name | Acc@1 | Acc@5 | #Params | SHA-256 (first 8 chars) | URL |
---|---|---|---|---|---|
CoaT-Lite Tiny | 77.5 | 93.8 | 5.7M | e88e96b0 | model, log |
CoaT-Lite Mini | 79.1 | 94.5 | 11M | 6b4a8ae5 | model, log |
CoaT-Lite Small | 81.9 | 95.5 | 20M | 8d362f48 | model, log |
CoaT-Lite Medium | 83.6 | 96.7 | 45M | a750cd63 | model, log |
CoaT Tiny | 78.3 | 94.0 | 5.5M | c6efc33c | model, log |
CoaT Mini | 81.0 | 95.2 | 10M | 40667eec | model, log |
CoaT Small | 82.1 | 96.1 | 22M | 7479cf9b | model, log |
The following commands provide an example (CoaT-Lite Tiny) to evaluate the pre-trained checkpoint.
# Download the pretrained checkpoint.
mkdir -p ./output/pretrained
wget http://vcl.ucsd.edu/coat/pretrained/coat_lite_tiny_e88e96b0.pth -P ./output/pretrained
sha256sum ./output/pretrained/coat_lite_tiny_e88e96b0.pth # Make sure it matches the SHA-256 hash (first 8 characters) in the table.
# Evaluate.
# Usage: bash ./scripts/eval.sh [model name] [output folder] [checkpoint path]
bash ./scripts/eval.sh coat_lite_tiny coat_lite_tiny_pretrained ./output/pretrained/coat_lite_tiny_e88e96b0.pth
# It should output results similar to "Acc@1 77.504 Acc@5 93.814" at very last.
The following commands provide an example (CoaT-Lite Tiny, 8-GPU) to train the CoaT model.
# Usage: bash ./scripts/train.sh [model name] [output folder]
bash ./scripts/train.sh coat_lite_tiny coat_lite_tiny
Note: Some training hyperparameters for CoaT Small and CoaT-Lite Medium are different from the default settings. We will update them soon.
The following commands provide an example (CoaT-Lite Tiny) to evaluate the checkpoint after training.
# Usage: bash ./scripts/eval.sh [model name] [output folder] [checkpoint path]
bash ./scripts/eval.sh coat_lite_tiny coat_lite_tiny_eval ./output/coat_lite_tiny/checkpoints/checkpoint0299.pth
@misc{xu2021coscale,
title={Co-Scale Conv-Attentional Image Transformers},
author={Weijian Xu and Yifan Xu and Tyler Chang and Zhuowen Tu},
year={2021},
eprint={2104.06399},
archivePrefix={arXiv},
primaryClass={cs.CV}
}
This repository is released under the Apache License 2.0. License can be found in LICENSE file.
Thanks to DeiT and pytorch-image-models for a clear and data-efficient implementation of ViT. Thanks to lucidrains' implementation of Lambda Networks and CPVT.