Hierarchical Pretraining: Research Repository

This is a research repository for reproducing the results from the project "Self-supervised pretraining improves self-supervised pretraining." You can find the arXiv prepint here: https://arxiv.org/abs/2103.12718

@article{reed2021self,
  title={Self-supervised pretraining improves self-supervised pretraining.},
  author={Reed, Colorado J and Yue, Xiangyu and Nrusimha, Ani and Ebrahimi, Sayna and Vijaykumar, Vivek and Mao, Richard and Li, Bo and Zhang, Shanghang and Guillory, Devin and Metzger, Sean and Keutzer, Kurt and Darrell, Trevor},
  journal={arXiv preprint arXiv:2103.12718},
  year={2021}
}

Pretrained Models

Note: we are currently adding many more pretrained models to this table.

ResNet-50

Pretraining Algorithm	Pretraining Data	Pretraining Amount	Type	Notes	Model URL
MoCo-v2	ImageNet	200 ep	Pytorch	67.7% ImageNet Top-1. codebase - paper	download
MoCo-v2-800	ImageNet	800 ep	Pytorch	71.1% ImageNet Top-1. codebase - paper	download
SimCLR-v1	ImageNet	800 ep	Tensorflow	69.1% ImageNet Top-1. codebase - paper	download
SimCLR-v2	ImageNet	800 ep	Tensorflow	71.7% ImageNet Top-1. codebase - paper	Google Storage (r50_1x_sk0/model.ckpt...)

Installation

# repo
git clone git@github.com:cjrd/base-training.git

# setup environment
conda create -n hpt python=3.7 ipython
conda activate hpt

# NOTE: if you are not using CUDA 10.2, you need to change the 10.2 in this file appropriately
# (check CUDA version with e.g. `cat /usr/local/cuda/version.txt`)
conda install pytorch==1.6.0 torchvision==0.7.0 cudatoolkit=10.2 -c pytorch

# install local submodules
cd OpenSelfSup
pip install -v -e .

Data installation

The data you need will depend on your goals, but generally speaking, download the RESISC-45 dataset to make sure you have everything working correctly.

Pretrained Models

cd OpenSelfSup/data/basetrain_chkpts/
./download-pretrained-models.sh

RESISC-45

RESISC-45 contains 31,500 aerial images, covering 45 scene classes with 700 images in each class.

# cd to the directory where you want the data, $DATA
wget https://people.eecs.berkeley.edu/~cjrd/data/resisc45.tar.gz
tar xf resisc45.tar.gz

# replace/set $DATA and $CODE as appropriate
ln -s $DATA/resisc45 $CODE/OpenSelfSup/data/resisc45/all

Verify Install

Check installation by pretraining using mocov2, extracting the model weights, evaluating the representations, and then viewing the results:

cd OpenSelfSup

# Sanity check: MoCo for 20 epoch on 4 gpus
./tools/dist_train.sh configs/selfsup/moco/r50_v2_resisc_in_basetrain_20ep.py 4

# make some variables so its clear what's happening
CHECKPOINT=work_dirs/selfsup/moco/r50_v2_resisc_in_basetrain_20ep/epoch_20.pth
BACKBONE=work_dirs/selfsup/moco/r50_v2_resisc_in_basetrain_20ep/epoch_20_moco_in_basetrain.pth
# Extract the backbone
python tools/extract_backbone_weights.py ${CHECKPOINT} ${BACKBONE}

# Evaluate the representations
./benchmarks/dist_train_linear.sh configs/benchmarks/linear_classification/resisc45/r50_last.py ${BACKBONE}

# View the results
cd work_dirs
# you may need to install tensorboard
tensorboard --logdir .

Base Training

Right now we assume ImageNet base trained models.

cd OpenSelfSup/data/basetrain_chkpts/
./download-pretrained-models.sh

Pretraining With a New Dataset

We have a handy set of config generators to make pretraining with a new dataset easy and consistent!

FIRST, you will need the image pixel mean/std of your dataset, if you don't have it, you can do:

./compute-dataset-pixel-mean-std.py --data /path/to/image-folder --numworkers 20 --batchsize 256

where image-folder has the structure from ImageFolder in pytorch
class/image-name.jp[e]g
or whatever image extension you're using

if your dataset is not arranged in this way, you can either: (i) use symlinks to put it in this structure (ii) update the above script to read in your data

NEXT, copy the pretraining template

cd utils
cp templates/pretraining-config-template.sh pretrain-configs/my-dataset-config.sh
# edit pretrain-configs/my-dataset-config.sh

# once edited, generate the project
./gen-pretrain-project.sh pretrain-configs/my-dataset-config.sh

What just happened? We generated a bunch of pretraining configs in the following location (take a loot at all of these files to get a feel for how this works):

OpenSelfSup/configs/hpt-pretrain/${shortname}

NEXT, you're ready to kick off a trial run to make sure the pretraining is working as expected =)

# the `-t` flag means `trial`: it'll only run 1 20 epoch pretraining
 ./utils/pretrain-runner.sh -t -d OpenSelfSup/configs/hpt-pretrain/${shortname}

NEXT, if this works, kick off the full training. NOTE: you can kick this off multiple times as long as the config directories share the same filesystem

# simply removing the `-t` flag from above
 ./utils/pretrain-runner.sh -d OpenSelfSup/configs/hpt-pretrain/${shortname}

NEXT, if you want to perform BYOL pretraining, add -b flag.

# simply add the `-b` flag to above. Currently, we only do it on Chexpert, Resisc, and Bdd for Exp3
 ./utils/pretrain-runner.sh -d OpenSelfSup/configs/hpt-pretrain/${shortname} -b

Congratulations: you've launch a full hierarchical pretraining experiment.

FAQs/PROBLEMS? * How does pretrain-runner.sh keep track of what's been pretrained? * In each config directory, it creates a .pretrain-status folder to keep track of what's processing/finished. See them with e.g. find OpenSelfSup/configs/hpt-pretrain -name '.pretrain-status' * How to redo a pretraining, e.g. because it crashed or something changed? Remove the * Remove the associate .proc or .done file. Find these e.g. bash find OpenSelfSup/configs/hpt-pretrain -name '.proc' find OpenSelfSup/configs/hpt-pretrain -name '.done'

Evaluating Pretrained Representations

This has been simplified to simply:

./utils/pretrain-evaluator.sh -b OpenSelfSup/work_dirs/hpt-pretrain/${shortname}/ -d OpenSelfSup/configs/hpt-pretrain/${shortname}

where -b is the backbone directory and -d is the config directory. This command also works for cross-dataset evaluation (e.g. evaluate models trained on Resic45 and evaluate on UC Merced dataset).

FAQ

Where are the checkpoints and logs? E.g., if you pass in configs/hpt-pretrain/resisc as the config directory, then the working directories for this evalution is e.g. work_dirs/hpt-pretrain/resisc/linear-eval/....

Finetuning

Assuming you generated the pretraining project as specified above, finetuning is as simple as:

./utils/finetune-runner.sh -d ./OpenSelfSup/configs/hpt-pretrain/${shortname}/finetune/ -b ./OpenSelfSup/work_dirs/hpt-pretrain/${shortname}/

where -b is the backbone directory and -d is the config directory Note: to finetune using other backbones, simply pass in a different backbone directory (the script searches for final_backbone.pth files in the provided directory tree)

Finetuning only on pretrained checkpoints with BEST linear analysis

First, specify the pretraining epochs which gives the best linear evaluation result in ./utils/top-linear-analysis-ckpts.txt. Here is an example:

# dataset best-moco-bt best-sup-bt best-no-bt
chest_xray_kids 5000 10000 100000
resisc 5000 50000 100000
chexpert 50000 50000 400000

, in which for chest_xray_kids dataset, 5000-iters, 10000-iters, 100000-iters are the best pretrained models under moco base-training, imagenet-supervised base-training, and no base-training, respectively.

Second, run the following command to perform finetuning only on the best checkpoints (same as above, except that the change of script name):

./utils/finetune-runner-top-only.sh -d ./OpenSelfSup/configs/hpt-pretrain/${shortname}/finetune/ -b ./OpenSelfSup/work_dirs/hpt-pretrain/${shortname}

Pretraining on top of pretraining

Using the output of previously pretrained models, it is very easy to correctly setup pretraining on top of the pretraining. Simply create a new config

utils/pretrain-configs/dataname1-dataname2.sh

(see resisc-ucmerced.sh for an example)

and then set the basetrained models to be the final_backbone.pth from the output of the last pretrained. e.g. for using resisc-45 outputs:

export basetrain_weights=(
    "work_dirs/hpt-pretrain/resisc/moco_v2_800ep_basetrain/50000-iters/final_backbone.pth"

    "work_dirs/hpt-pretrain/resisc/imagenet_r50_supervised_basetrain/50000-iters/final_backbone.pth"

    "work_dirs/hpt-pretrain/resisc/no_basetrain/200000-iters/final_backbone.pth"
)

(see resisc-ucmerced.sh for an example)

To select which backbones to use, evaluate the linear performance from the various source outputs (e.g. all the resisc pretrained outputs) on the target data (e.g. on uc-merced data).

Then simply generate the project and execute the pretraining as normal:

./gen-pretrain-project.sh pretrain-configs/dataname1-dataname2.sh

./pretrain-runner.sh -d OpenSelfSup/configs/hpt-pretrain/$dataname1-dataname2

Object Detection / Semantic Segmentation

Object detection/segmentation uses detectron2 and takes place in the directory

OpenSelfSup/benchmarks/detection

First: Check if the dataset configs you need are already present in configs. E.g. if you're working with CoCo, you'll see the following 2 configs:

configs/coco_R_50_C4_2x.yaml
configs/coco_R_50_C4_2x_moco.yaml

We'll use the config with the _moco suffix for all obj det and segmentation. If your configs already exist, skip the next step.

Next: assuming your configs do not exist, set up the configs you need for your dataset by copying an existing set of configs

cp configs/coco_R_50_C4_2x.yaml ${MYDATA}_R50_C4_2x.yaml
cp configs/coco_R_50_C4_2x_moco.yaml ${MYDATA}_R50_C4_2x_moco.yaml

Edit ${MYDATA}_R50_C4_2x.yaml and set MIN_SIZE_TRAIN and MIN_SIZE_TEST to be appropriate for your dataset. Also, rename TRAIN and TEST to have your dataset name, set MASK_ON to True if doing semantic segmentation, and update STEPS and MAX_ITER if running the training for a different amount of time is appropriate (check relevant publications / codebases to set the training schedule).

Edit ${MYDATA}_R50_C4_2x_moco.yaml and set PIXEL_MEAN and PIXEL_STD (use compute-dataset-pixel-mean-std.py script above, if you don't know them).

Then, edit train_net.py and add the appropriate data registry lines for your train/val data

register_coco_instances("dataname_train", {}, "obj-labels-in-coco-format_train.json", "datasets/dataname/dataname_train")
register_coco_instances("dataname_val", {}, "obj-labels-in-coco-format_val.json", "datasets/dataname/dataname_val")

Then, setup symlinks to your data under datasets/dataname/dataname_train and datasets/dataname/dataname_val, where you replace dataname with your dataname used in the config/registry.

Next, convert your backbone(s) to detectron format, e.g. (NOTE: I recommend keeping backbones in the same directory that they are originally present in, and appending a -detectron2 suffix)

python convert-pretrain-to-detectron2.py ../../data/basetrain_chkpts/imagenet_r50_supervised.pth ../../data/basetrain_chkpts/imagenet_r50_supervised-detectron2.pth

Next kick off training

python train_net.py --config-file configs/DATANAME_R_50_C4_24k_moco.yaml --num-gpus 4 OUTPUT_DIR results/${UNIQUE_DATANAME_EXACTLY_DESCRIBING_THIS_RUN}/ TEST.EVAL_PERIOD 2000 MODEL.WEIGHTS ../../data/basetrain_chkpts/imagenet_r50_supervised-detectron2.pth SOLVER.CHECKPOINT_PERIOD ${INT_HOW_OFTEN_TO_CHECKPOINT}

results will be in results/${UNIQUE_DATANAME_EXACTLY_DESCRIBING_THIS_RUN}, and you can use tensorboard to view them.

Commit and Share Results

Run the following command to grab all results (linear analysis, finetunes, etc) and put them into the appropriate json results file in results/:

./utils/update-all-results.sh

You can verify the