News: The official TF2 re-implementation is available in DeepLab2. Axial-SWideRNet achieves 68.0% PQ or 83.5% mIoU on Cityscaspes validation set, with only single-scale inference and ImageNet-1K pretrained checkpoints.
This is a PyTorch re-implementation of the Axial-DeepLab paper. The re-implementation is mainly done by an amazing senior student, Huaijin Pi.
@inproceedings{wang2020axial,
title={Axial-DeepLab: Stand-Alone Axial-Attention for Panoptic Segmentation},
author={Wang, Huiyu and Zhu, Yukun and Green, Bradley and Adam, Hartwig and Yuille, Alan and Chen, Liang-Chieh},
booktitle={European Conference on Computer Vision (ECCV)},
year={2020}
}Currently, only ImageNet classification with the "Conv-Stem + Axial-Attention" backbone is supported. If you are interested in contributing to this repo, please open an issue and we can further discuss.
pip install tensorboardX
mkdir data
cd data
ln -s path/to/dataset imagenet- Non-distributed training
python train.py --model axial50s --gpu_id 0,1,2,3 --batch_size 128 --val_batch_size 128 --name axial50s --lr 0.05 --nesterov- Distributed training
CUDA_VISIBLE_DEVICES=0,1,2,3 python dist_train.py --model axial50s --batch_size 128 --val_batch_size 128 --name axial50s --lr 0.05 --nesterov --dist-url 'tcp://127.0.0.1:4128' --dist-backend 'nccl' --multiprocessing-distributed --world-size 1 --rank 0You can change the model name to train different models.
python train.py --model axial50s --gpu_id 0,1,2,3 --batch_size 128 --val_batch_size 128 --name axial50s --lr 0.05 --nesterov --testYou can test with distributed settings in the same way.
| Method | Params (M) | Top-1 Acc (%) |
|---|---|---|
| ResNet-26 | 13.7 | 74.5 |
| Axial-ResNet-26-S | 5.9 | 75.8 |
-
ImageNet training script is modified from https://github.com/mit-han-lab/proxylessnas
-
ImageNet distributed training script is modified from https://github.com/pytorch/examples/tree/master/imagenet
-
ResNet is modified from https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py