AutoFocusFormer

AFF-Base:

This software project accompanies the research paper, AutoFocusFormer: Image Segmentation off the Grid (CVPR 2023).

Chen Ziwen, Kaushik Patnaik, Shuangfei Zhai, Alvin Wan, Zhile Ren, Alex Schwing, Alex Colburn, Li Fuxin

arXiv | video narration | AFF-Classification (this repo) | AFF-Segmentation

Introduction

AutoFocusFormer (AFF) is the first adaptive-downsampling network capable of dense prediction tasks such as semantic/instance segmentation.

AFF abandons the traditional grid structure of image feature maps, and automatically learns to retain the most important pixels with respect to the task goal.

AFF consists of a local-attention transformer backbone and a task-specific head. The backbone consists of four stages, each stage containing three modules: balanced clustering, local-attention transformer blocks, and adaptive downsampling.

AFF demonstrates significant savings on FLOPs (see our models with 1/5 downsampling rate), and significant improvement on recognition of small objects.

Notably, AFF-Small achieves 44.0 instance segmentation AP and 66.9 panoptic segmentation PQ on Cityscapes val with a backbone of only 42.6M parameters, a performance on par with Swin-Large, a backbone with 197M params (saving 78%!).

Main Results on ImageNet with Pretrained Models

FPS is obtained on a single V100 GPU.

We train with a total batch size 4096.

Getting Started

Clone this repo

git clone [email protected]:apple/ml-autofocusformer.git
cd ml-autofocusformer

One can download pre-trained checkpoints through the links in the table above.

Create environment and install requirements

sh create_env.sh

See further documentation inside the script file.

Our experiments are run with CUDA==11.6 and pytorch==1.12.

Prepare data

We use standard ImageNet dataset, which can be downloaded from http://image-net.org/.

For standard folder dataset, move validation images to labeled sub-folders. The file structure should look like:

$ tree imagenet
imagenet/
 training
    class1
       img1.jpeg
       img2.jpeg
       ...
    class2
       img3.jpeg
       ...
    ...
 validation
     class1
        img4.jpeg
        img5.jpeg
        ...
     class2
        img6.jpeg
        ...
     ...

Train and evaluate

Modify the arguments in script run_aff.sh (e.g., path to dataset) and run

sh run_aff.sh

for training or evaluation.

Run python main.py -h to see full documentation of the args.

One can also directly modify the config files in configs/.

Citing AutoFocusFormer

@inproceedings{autofocusformer,
    title = {AutoFocusFormer: Image Segmentation off the Grid},
    booktitle = {Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR)},
    author = {Ziwen, Chen and Patnaik, Kaushik and Zhai, Shuangfei and Wan, Alvin and Ren, Zhile and Schwing, Alex and Colburn, Alex and Fuxin, Li},
    year = {2023},
}