Environment Setup

To set up your environment, execute the following commands:

conda create -n rein -y
conda activate rein
conda install pytorch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 pytorch-cuda=11.7 -c pytorch -c nvidia -y
pip install -U openmim
mim install mmengine
mim install "mmcv>=2.0.0"
pip install "mmsegmentation>=1.0.0"
pip install "mmdet>=3.0.0"
pip install xformers=='0.0.20' # optional for DINOv2
pip install -r requirements.txt
pip install future tensorboard

Pretraining Weights

• Download: Download pre-trained weights from facebookresearch for testing. Place them in the project directory without changing the file name.
• Convert: Convert pre-trained weights for training or evaluation.

  python tools/convert_models/convert_dinov2.py dinov2_vitl14_pretrain.pth checkpoints/dinov2_converted.pth

  (optional for 1024x1024 resolution)

  python tools/convert_models/convert_dinov2.py dinov2_vitl14_pretrain.pth checkpoints/dinov2_converted_1024x1024.pth --height 1024 --width 1024

Training

Start training in single GPU:

python tools/train.py configs/dinov2/rein_dinov2_mask2former_512x512_bs1x4.py

Start training in multiple GPU:

PORT=12345 CUDA_VISIBLE_DEVICES=1,2,3,4 bash tools/dist_train.sh configs/dinov2/rein_dinov2_mask2former_1024x1024_bs4x2.py NUM_GPUS

Inference

python tools/visualize.py configs/dinov2/rein_dinov2_mask2former_512x512_bs1x4.py /PATH/TO/FINAL/CHECKPOINT/iter_40000.pth /INFERENCE/DATASET/images --suffix ".jpg" --backbone "/PATH/TO/FACEBOOK/BACKBONE/dinov2_converted.pth" --save_dir "/DIRECTORY/TO/STORE/PREDICTIONS"

Generate full weights

Because we only fine-tune and save the REIN and head weights, if you need a complete set of segmentor weights, you need to use this script:

python generate_full_weights.py --segmentor_save_path SEGMENTOR_SAVE_PATH --backbone CONVERTED_BACKBONE --rein_head REIN_HEAD