Paper Links: Our most recent TPAMI version with improvements and extensions (Earlier ICCV version).
By Zilong Huang, Xinggang Wang, Yunchao Wei, Lichao Huang, Chang Huang, Humphrey Shi, Wenyu Liu and Thomas S. Huang.
2020/07: Stay tuned for our newest TPAMI version of code with improvements (CityScapes test: 81.9%, ADE20K val: 45.76%, LIP Val: 55.47%) and extensions (3D-CCNet on video dataset CamVid test: 79.1%). Code to be released soon!
2019/08: The new version CCNet is released on branch Pytorch-1.1 which supports Pytorch 1.0 or later and distributed multiprocessing training and testing This current code is a implementation of the experiments on Cityscapes in the CCNet ICCV version. We implement our method based on open source pytorch segmentation toolbox.
2018/12: Renew the code and release trained models with R=1,2. The trained model with R=2 achieves 79.74% on val set and 79.01% on test set with single scale testing.
2018/11: Code released.
Long-range dependencies can capture useful contextual information to benefit visual understanding problems. In this work, we propose a Criss-Cross Network (CCNet) for obtaining such important information through a more effective and efficient way. Concretely, for each pixel, our CCNet can harvest the contextual information of its surrounding pixels on the criss-cross path through a novel criss-cross attention module. By taking a further recurrent operation, each pixel can finally capture the long-range dependencies from all pixels. Overall, our CCNet is with the following merits:
- GPU memory friendly
- High computational efficiency
- The state-of-the-art performance
Overview of the proposed CCNet for semantic segmentation. The proposed recurrent criss-cross attention takes as input feature maps H and output feature maps H'' which obtain rich and dense contextual information from all pixels. Recurrent criss-cross attention module can be unrolled into R=2 loops, in which all Criss-Cross Attention modules share parameters.
To get a deeper understanding of our RCCA, we visualize the learned attention masks as shown in the figure. For each input image, we select one point (green cross) and show its corresponding attention maps when R=1 and R=2 in columns 2 and 3 respectively. In the figure, only contextual information from the criss-cross path of the target point is capture when R=1. By adopting one more criss-cross module, ie, R=2 the RCCA can finally aggregate denser and richer contextual information compared with that of R=1. Besides, we observe that the attention module could capture semantic similarity and long-range dependencies.
CCNet is released under the MIT License (refer to the LICENSE file for details).
If you find CCNet useful in your research, please consider citing:
@article{huang2020ccnet,
author={Huang, Zilong and Wang, Xinggang and Wei, Yunchao and Huang, Lichao and Shi, Humphrey and Liu, Wenyu and Huang, Thomas S.},
journal={IEEE Transactions on Pattern Analysis and Machine Intelligence},
title={CCNet: Criss-Cross Attention for Semantic Segmentation},
year={2020},
month={},
volume={},
number={},
pages={1-1},
keywords={Semantic Segmentation;Graph Attention;Criss-Cross Network;Context Modeling},
doi={10.1109/TPAMI.2020.3007032},
ISSN={1939-3539}}
@article{huang2018ccnet,
title={CCNet: Criss-Cross Attention for Semantic Segmentation},
author={Huang, Zilong and Wang, Xinggang and Huang, Lichao and Huang, Chang and Wei, Yunchao and Liu, Wenyu},
booktitle={ICCV},
year={2019}}
To install PyTorch==0.4.0 or 0.4.1, please refer to https://github.com/pytorch/pytorch#installation.
4 x 12G GPUs (e.g. TITAN XP)
Python 3.6
gcc (GCC) 4.8.5
CUDA 8.0
Some parts of InPlace-ABN and Criss-Cross Attention have native CUDA implementations, which must be compiled with the following commands:
cd libs
sh build.sh
python build.py
cd ../cc_attention
sh build.sh
python build.py
The build.sh
script assumes that the nvcc
compiler is available in the current system search path.
The CUDA kernels are compiled for sm_50
, sm_52
and sm_61
by default.
To change this (e.g. if you are using a Kepler GPU), please edit the CUDA_GENCODE
variable in build.sh
.
Plesae download cityscapes dataset and unzip the dataset into YOUR_CS_PATH
.
Please download MIT imagenet pretrained resnet101-imagenet.pth, and put it into dataset
folder.
Training script.
python train.py --data-dir ${YOUR_CS_PATH} --random-mirror --random-scale --restore-from ./dataset/resnet101-imagenet.pth --gpu 0,1,2,3 --learning-rate 1e-2 --input-size 769,769 --weight-decay 1e-4 --batch-size 8 --num-steps 60000 --recurrence 2
【Recommend】You can also open the OHEM flag to reduce the performance gap between val and test set.
python train.py --data-dir ${YOUR_CS_PATH} --random-mirror --random-scale --restore-from ./dataset/resnet101-imagenet.pth --gpu 0,1,2,3 --learning-rate 1e-2 --input-size 769,769 --weight-decay 1e-4 --batch-size 8 --num-steps 60000 --recurrence 2 --ohem 1 --ohem-thres 0.7 --ohem-keep 100000
Evaluation script.
python evaluate.py --data-dir ${YOUR_CS_PATH} --restore-from snapshots/CS_scenes_60000.pth --gpu 0 --recurrence 2
All in one.
./run_local.sh YOUR_CS_PATH
We run CCNet with R=1,2 three times on cityscape dataset separately and report the results in the following table. Please note there exist some problems about the validation/testing set accuracy gap (1~2%). You need to run multiple times to achieve a small gap or turn on OHEM flag. Turning on OHEM flag also can improve the performance on the val set. In general, I recommend you use OHEM in training step.
We train all the models on fine training set and use the single scale for testing. The trained model with R=2 79.74 can also achieve about 79.01 mIOU on cityscape test set with single scale testing (for saving time, we use the whole image as input).
R | mIOU on cityscape val set (single scale) | Link |
---|---|---|
1 | 77.31 & 77.91 & 76.89 | 77.91 |
2 | 79.74 & 79.22 & 78.40 | 79.74 |
2+OHEM | 78.67 & 80.00 & 79.83 | 80.00 |
We thank NSFC, ARC DECRA DE190101315, ARC DP200100938, HUST-Horizon Computer Vision ResearchCenter, and IBM-ILLINOIS Center for Cognitive ComputingSystems Research (C3SR).
Self-attention related methods:
Object Context Network
Dual Attention Network
Semantic segmentation toolboxs:
pytorch-segmentation-toolbox
semantic-segmentation-pytorch
PyTorch-Encoding