Skip to content

Latest commit

 

History

History
 
 

pssl

README.md

PSSL: Distilling ensemble of explanations for weakly-supervised pre-training of image segmentation models.

Reference

Xuhong Li, Haoyi Xiong, Yi Liu, Dingfu Zhou, Zeyu Chen, Yaqing Wang, and Dejing Dou. "Distilling ensemble of explanations for weakly-supervised pre-training of image segmentation models." Machine Learning (2022): 1-17. https://arxiv.org/abs/2207.03335

Overview

In our work, we propose a method that leverages Pseudo Semantic Segmentation Labels (PSSL), to enable the end-to-end pre-training for image segmentation models based on classification datasets. PSSL was inspired by the observation that the explanation results of classification models, obtained through explanation algorithms such as CAM, SmoothGrad and LIME, would be close to the pixel clusters of visual objects. Specifically, PSSL is obtained for each image by interpreting the classification results and aggregating an ensemble of explanations queried from multiple classifiers to lower the bias caused by single models. With PSSL for every image of ImageNet, the proposed method leverages a weighted segmentation learning procedure to pre-train the segmentation network en masse.

arch

Training

Our method improves the pre-training step of segmentation models. We provide the details of the pre-training here, but note that the pre-training step can be skipped because we have provided the pre-trained models. The PSSL dataset can be obtained by sending an email to [email protected] via an official email (not use qq, gmail, etc.) including your institution/company information and the purpose on the dataset.

Here we show the configuration files of two lightweight models, STDC2 and PPLite-Seg-B, where the download links are also provided.

(Optional) Pretraining

There is no need to do this step if STDC2 and PPLite-Seg-B are used.

Otherwise, checking the following steps for preparing the enviroment and the datasets:

  1. Make sure that PaddleSeg is well installed. See the installation guide for details.

  2. Make sure that ImageNet has been downloaded (~138G) and extracted to data/ImageNet_org. We only use the training set of ImageNet here. See below the dataset structure.

  3. After getting the PSSL dataset link, download and extract to data/pssl2.1_consensus. See below the dataset structure.

Make sure that the datasets have structures as follows:

PaddleSeg
│   ...  
│
└───data
│   │
│   └───ImageNet_org
│   │   │  
│   │   └───train
│   │       │  
│   │       └───n01440764
│   │       │   │   n01440764_10026.JPEG
│   │       │   │   ...
│   │       │  
│   │       └───nxxxxxxxx
│   │  
│   └───pssl2.1_consensus  
│   │   │   imagenet_lsvrc_2015_synsets.txt
│   │   │   train.txt
│   │   └───train
│   │       │  
│   │       └───n01440764
│   │       │   │   n01440764_10026.JPEG_eiseg.npz
│   │       │   │   ...
│   │       │  
│   │       └───nxxxxxxxx
│   │       │   ...

Having installed PaddlePaddle and PaddleSeg and prepared datasets (ImageNet and PSSL), we can run the pre-training script:

  • For STDC-Seg
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7

path_save="work_dirs_stdc2_pssl"
python -m paddle.distributed.launch --log_dir $path_save tools/train.py \
       --config configs/pssl/stdc2_seg_pssl.yml \
       --log_iters 200 \
       --num_workers 12 \
       --save_interval 13345 \
       --keep_checkpoint_max 20 \
       --save_dir ${path_save}/snapshot
  • For PP-LiteSeg
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7

path_save="work_dirs_pp_liteseg_stdc2_pssl"
python -m paddle.distributed.launch --log_dir $path_save tools/train.py \
       --config configs/pssl/pp_liteseg_stdc2_pssl.yml \
       --log_iters 100 \
       --num_workers 12 \
       --save_interval 13345 \
       --keep_checkpoint_max 20 \
       --save_dir ${path_save}/snapshot

After the pre-training, the weights are saved in ${path_save}/snapshot/iter_xxx/model.pdparams. Conventionall, we use the 5th epoch's checkpoint, i.e., 66725th iter (total batch size = 12 * 8 = 96), to do the downstream tasks.

For other models, modify the configuration file as needed.

Downstream Tasks

For downstream tasks, PSSL does not need to change anything, except loading the pre-trained model. We can change this by simply adding one line in the config file, for example,

model:
  ...
  pretrained: work_dirs_pp_liteseg_stdc2_pssl/snapshot/iter_66725/model.pdparams

Performance

Pascal VOC 2012 + Aug

Model Backbone Resolution Training Iters mIoU mIoU (flip) mIoU (ms+flip) Links
STDC2-Seg50 STDC2 512x512 40000 68.98% 70.07% 69.99% model | log | vdl
STDC2-Seg50 + PSSL STDC2 512x512 40000 74.49% 74.96% 75.79% model | log

Cityscapes

Model Backbone Training Iters Train Crops Test Resolution mIoU mIoU (flip) mIoU (ms+flip) Links
PP-LiteSeg-B STDC2 160000 1024x512 2048x1024 79.04% 79.52% 79.85% config|model|log|vdl
PP-LiteSeg-B + PSSL STDC2 160000 1024x512 2048x1024 79.06% 79.61% 79.97% config|model|log