train.md

Train

Training Dataset

You should prepare Adobe Image Matting Dataset and COCO for training.

• Get Adobe Image Matting Dateset.
• For training, merge 'Adobe-licensed images' and 'Other' folder to use all 431 foregrounds and alphas.
• Get COCO.

Check ViTMatte/configs/common/dataloader to modify PATH TO TRAINING DATA.

Pretrained Model

• Get DINO pretrained ViT-S and MAE pretrained ViT-B.

• Or you could download and preprocess pretrained weights by

  cd ViTMatte/pretrained
  python preprocess.py
  

Train ViTMatte

ViTMatte has 2 sizes: ViTMatte-S and ViTMatte-B.

You can modify configs in configs/ViTMatte_S_100ep.py and configs/ViTMatte_B_100ep.py.

Run:

    python main.py \
        --config-file configs/ViTMatte_S_100ep.py \
        --num-gpus 2

to train ViTMatte-S.

Run:

    python main.py \
        --config-file configs/ViTMatte_B_100ep.py \
        --num-gpus 2

to train ViTMatte-B.

Training output will be saved in ViTMatte/output_of_train/