This repo contains the supported code and configuration files to reproduce semantic segmentaion results of AS-MLP. It is based on Swin Transformer.
Backbone | Method | Crop Size | Lr Schd | mIoU (ms+flip) | Params | FLOPs | config | model |
---|---|---|---|---|---|---|---|---|
AS-MLP-T | UPerNet | 512x512 | 160K | 46.5 | 60M | 937G | config | |
AS-MLP-S | UperNet | 512x512 | 160K | 49.2 | 81M | 1024G | config | |
AS-MLP-B | UperNet | 512x512 | 160K | 49.5 | 121M | 1166G | config | onedrive |
Notes:
- Pre-trained models can be downloaded from AS-MLP for ImageNet Classification.
Please refer to get_started.md for installation and dataset preparation.
# single-gpu testing
python tools/test.py <CONFIG_FILE> <SEG_CHECKPOINT_FILE> --eval mIoU
# multi-gpu testing
tools/dist_test.sh <CONFIG_FILE> <SEG_CHECKPOINT_FILE> <GPU_NUM> --eval mIoU
# multi-gpu, multi-scale testing
tools/dist_test.sh <CONFIG_FILE> <SEG_CHECKPOINT_FILE> <GPU_NUM> --aug-test --eval mIoU
To train with pre-trained models, run:
# single-gpu training
python tools/train.py <CONFIG_FILE> --options model.pretrained=<PRETRAIN_MODEL> [model.backbone.use_checkpoint=True] [other optional arguments]
# multi-gpu training
tools/dist_train.sh <CONFIG_FILE> <GPU_NUM> --options model.pretrained=<PRETRAIN_MODEL> [model.backbone.use_checkpoint=True] [other optional arguments]
For example, to train an UPerNet model with a AS-MLP-T
backbone and 8 gpus, run:
tools/dist_train.sh configs/asmlp/upernet_asmlp_tiny_patch4_shift5_512x512_160k_ade20k.py 8 --options model.pretrained=<PRETRAIN_MODEL>
Notes:
use_checkpoint
is used to save GPU memory. Please refer to this page for more details.- The default learning rate and training schedule is for 8 GPUs and 2 imgs/gpu.
@article{Lian_2021_ASMLP,
author = {Lian, Dongze and Yu, Zehao and Sun, Xing and Gao, Shenghua},
title = {AS-MLP: An Axial Shifted MLP Architecture for Vision},
journal={ICLR},
year = {2022}
}
Image Classification: See AS-MLP for Image Classification.
Object Detection: See AS-MLP for Object Detection.