Skip to content

Latest commit

 

History

History
358 lines (253 loc) · 30.9 KB

Readme.md

File metadata and controls

358 lines (253 loc) · 30.9 KB
Raw

Advancing Plain Vision Transformer Towards Remote Sensing Foundation Model

Di Wang, Qiming Zhang, Yufei Xu, Jing Zhang, Bo Du, Dacheng Tao and Liangpei Zhang

Updates | Introduction | Results & Models | Usage |

PWC PWC PWC PWC PWC PWC PWC PWC PWC PWC

Current applications

ViTAE: Please see ViTAE-Transformer;

VSA: Please see ViTAE-VSA;

Matting: Please see ViTAE-Transformer for matting;

Remote Sensing Pretraining: Please see ViTAE-Transformer-Remote-Sensing;

Updates

2023.10.18

RVSA won the highly cited paper!

2023.03.18

ViTAE-B + RVSA helped us win the championship of "High Resolution SAR Image Coastal Aquaculture Farm Segmentation Track" in "The 5th Gaofen Challenge", Team: TNT. （第五届“中科星图杯”国际高分遥感图像解译大赛高分辨率SAR图像中近海养殖场分割赛道冠军）News

2023.01.18

Our models have been supported by LuoJiaNET, please refer to RS-Vision-Foundation-Models for more details.

2022.11.21

The early access is available! TGRS link

2022.11.15

The arXiv has been updated! arXiv link

2022.11.06

The paper has been accepted by IEEE TGRS!

2022.10.11

The codes, configs and training logs of segmentation in fintuning are released!

2022.10.09

The codes, configs and training logs of detection in fintuning are released!

2022.10.08

The codes of pretraining and classification in fintuning are released!

2022.09.19

The codes and training logs of the VSA have been released, which is the foundation of our RVSA.

Introduction

This repository contains codes, models and test results for the paper "Advancing Plain Vision Transformer Towards Remote Sensing Foundation Model".

We resort to plain vision transformers with about 100M and make the first attempt to propose large vision models customized for RS tasks and propose a new rotated varied-size window attention (RVSA) to substitute the original full attention to handle the large image size and objects of various orientations in RS images. The RVSA could significantly reduce the computational cost and memory footprint while learn better object representation by extracting rich context from the generated diverse windows.

Fig.1 - The pipeline of pretraining and finetuning.

 

Fig.2 - The structure and block of the adopted plain vision transformer, and the proposed RVSA.

Results and Models

Pretraining

MillionAID

Pretrain Backbone Input size Params (M) Pretrained model
MAE ViT-B 224 × 224 86 Weights
MAE ViTAE-B 224 × 224 89 Weights

Object Detection

DOTA-V1.0 Single-Scale

Method Pretrain Backbone Lr schd mAP Config Log Model
Oriented R-CNN MAE ViT-B + RVSA 1x 78.75 Config Log Model
Oriented R-CNN MAE ViT-B + RVSA $^ \Diamond$ 1x 78.61 Config Log Model
Oriented R-CNN MAE ViTAE-B + RVSA 1x 78.96 Config Log Model
Oriented R-CNN MAE ViTAE-B + RVSA $^ \Diamond$ 1x 78.99 Config Log Model

DOTA-V1.0 Multi-Scale

Method Pretrain Backbone Lr schd mAP Config Log Model
Oriented R-CNN MAE ViT-B + RVSA 1x 81.01 Config Log Model
Oriented R-CNN MAE ViT-B + RVSA $^ \Diamond$ 1x 80.80 Config Log Model
Oriented R-CNN MAE ViTAE-B + RVSA 1x 81.24 Config Log Model
Oriented R-CNN MAE ViTAE-B + RVSA $^ \Diamond$ 1x 81.18 Config Log Model

DIOR-R

Method Pretrain Backbone Lr schd mAP Config Log Model
Oriented R-CNN MAE ViT-B + RVSA 1x 70.67 Config Log Model
Oriented R-CNN MAE ViT-B + RVSA $^ \Diamond$ 1x 70.85 Config Log Model
Oriented R-CNN MAE ViTAE-B + RVSA 1x 70.95 Config Log Model
Oriented R-CNN MAE ViTAE-B + RVSA $^ \Diamond$ 1x 71.05 Config Log Model

Scene Classification

Pretrain Backbone UCM-55 AID-28 AID-55 NWPU-19 NWPU-28
MAE ViT-B + RVSA 99.70 96.92 98.33 93.79 95.49
Model Model Model Model Model
MAE ViT-B + RVSA $^ \Diamond$ 99.58 96.86 98.44 93.74 95.45
Model Model Model Model Model
MAE ViTAE-B + RVSA 99.56 97.03 98.48 93.93 95.69
Model Model Model Model Model
MAE ViTAE-B + RVSA $^ \Diamond$ 99.50 97.01 98.50 93.92 95.66
Model Model Model Model Model

Semantic Segmentation

ISPRS Potsdam

Method Pretrain Backbone Crop size Lr schd OA Config Log Model
UperNet MAE ViT-B + RVSA 512 × 512 160k 90.60 Config Log Model
UperNet MAE ViT-B + RVSA $^ \Diamond$ 512 × 512 160k 90.77 Config Log Model
UperNet MAE ViTAE-B + RVSA 512 × 512 160k 91.22 Config Log Model
UperNet MAE ViTAE-B + RVSA $^ \Diamond$ 512 × 512 160k 91.15 Config Log Model

iSAID

Method Pretrain Backbone Crop size Lr schd mIOU Config Log Model
UperNet MAE ViT-B + RVSA 896 × 896 160k 63.76 Config Log Model
UperNet MAE ViT-B + RVSA $^ \Diamond$ 896 × 896 160k 63.85 Config Log Model
UperNet MAE ViTAE-B + RVSA 896 × 896 160k 63.48 Config Log Model
UperNet MAE ViTAE-B + RVSA $^ \Diamond$ 896 × 896 160k 64.49 Config Log Model

LoveDA

Method Pretrain Backbone Crop size Lr schd mIOU Config Log Model
UperNet MAE ViT-B + RVSA 512 × 512 160k 51.95 Config Log Model
UperNet MAE ViT-B + RVSA $^ \Diamond$ 512 × 512 160k 51.95 Config Log Model
UperNet MAE ViTAE-B + RVSA 512 × 512 160k 52.26 Config Log Model
UperNet MAE ViTAE-B + RVSA $^ \Diamond$ 512 × 512 160k 52.44 Config Log Model

Usage

Environment:

  • Python 3.8.5
  • Pytorch 1.9.0+cu111
  • torchvision 0.10.0+cu111
  • timm 0.4.12
  • mmcv-full 1.4.1

Pretraining & Finetuning-Classification

Pretraining (8 × A100 GPUs, 3~5 days)

  1. Preparing the MillionAID: Download the MillionAID. Here, we use previous train_labels.txt and valid_labels.txt of the RSP, which contain labels. However, since we conduct the unsupervised pretraining, the labels are not necessary. It is easy for users to record image names and revise corresponding codes MAEPretrain_SceneClassification/util/datasets.py/class MillionAIDDataset.

  2. Pretraining: take ViT-B as an example (batchsize: 2048=8*256)

python -m torch.distributed.launch --nproc_per_node 8 --master_port 10000 main_pretrain.py \
--dataset 'millionAID' --model 'mae_vit_base_patch16' \
--batch_size 256 --epochs 1600 --warmup_epochs 40 \
--input_size 224 --mask_ratio 0.75 \
--blr 1.5e-4  --weight_decay 0.05 --gpu_num 8 \
--output_dir '../mae-main/output/'

Note: Padding the convolutional kernel of PCM in the pretrained ViTAE-B with convertK1toK3.py for finetuning.

  1. Linear probe: an example of evaluating the pretrained ViT-B on UCM-55
CUDA_VISIBLE_DEVICES=0 python -m torch.distributed.launch --nproc_per_node 1 --master_port 10000 main_linprobe.py \
--dataset 'ucm' --model 'vit_base_patch16' \
--batch_size 256 --epochs 100 --warmup_epochs 10 \
--blr 1e-1  --weight_decay 0 --tag 0 \
--finetune '../mae-main/output/millionAID_224/1600_0.75_0.00015_0.05_2048/checkpoint-1599.pth'

Finetuning evaluation for pretraining & Finetuning-Classification

For instance, finetuning ViTAE-B + RVSA on NWPU-28

CUDA_VISIBLE_DEVICES=0 python -m torch.distributed.launch --nproc_per_node 1 --master_port 20000 main_finetune.py \
--dataset 'nwpu' --model 'vitae_nc_base_win_rvsa' --input_size 224 --postfix 'sota' \
--batch_size 64 --epochs 200 --warmup_epochs 5 \
--blr 1e-3  --weight_decay 0.05 --split 28 --tag 0 --exp_num 1 \
--finetune '../mae-main/output/mae_vitae_base_pretrn/millionAID_224/1600_0.75_0.00015_0.05_2048/checkpoint-1599-transform-no-average.pth'

Finetuning-Detection & Finetuning-Segmentation

Since we use OBBDetection and MMSegmenation to implement corresponding detection or segmentation models, we only provide necessary config and backbone files. The main frameworks are both in RSP

git clone https://github.com/ViTAE-Transformer/ViTAE-Transformer-Remote-Sensing.git

The installation and dataset preparation can separately refer OBBDetection-installation and MMSegmentation-installation

Then put these files into corresponding folders.

For convenience, we preserve the relative path for users to find files.

For example, put ./Object Detection/mmdet/models/backbones/vit_win_rvsa_v3_wsz7.py into ViTAE-Transformer-Remote-Sensing/Object Detection/mmdet/models/backbones

Training-Detection

First, cd ./Object Detection

Then, we provide several examples. For instance,

Training the Oriented-RCNN with ViT-B + RVSA on DOTA-V1.0 multi-scale detection dataset with 2 GPUs

CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 --master_port=40000 tools/train.py \
configs/obb/oriented_rcnn/vit_base_win/faster_rcnn_orpn_our_rsp_vit-base-win-rvsa_v3_wsz7_fpn_1x_dota10_ms_lr1e-4_ldr75_dpr15.py \
--launcher 'pytorch' --options 'find_unused_parameters'=True

Training the Oriented-RCNN with ViTAE-B + RVSA $^ \Diamond$ backbone on DIOR-R detection dataset with 1 GPU

CUDA_VISIBLE_DEVICES=0 python -m torch.distributed.launch --nproc_per_node=1 --master_port=40001 tools/train.py \
configs/obb/oriented_rcnn/vit_base_win/faster_rcnn_orpn_our_rsp_vitae-nc-base-win-rvsa_v3_kvdiff_wsz7_fpn_1x_dior_lr1e-4_ldr75_dpr10.py \
--launcher 'pytorch' --options 'find_unused_parameters'=True

Inference-Detection

Predicting the saving detection map using ViT-B + RVSA $^ \Diamond$ on DOTA-V1.0 scale-scale detection dataset

CUDA_VISIBLE_DEVICES=0 python tools/test.py configs/obb/oriented_rcnn/vit_base_win/faster_rcnn_orpn_our_rsp_vit-base-win-rvsa_v3_kvdiff_wsz7_fpn_1x_dota10_lr1e-4_ldr75_dpr15.py \
../OBBDetection/work_dirs/faster/faster_rcnn_orpn_our_rsp_vit-base-win-rvsa_v3_kvdiff_wsz7_fpn_1x_dota10_lr1e-4_ldr75_dpr15/latest.pth \
--format-only --show-dir work_dirs/save/faster/display/faster_rcnn_orpn_our_rsp_vit-base-win-rvsa_v3_kvdiff_wsz7_fpn_1x_dota10_lr1e-4_ldr75_dpr15 \
--options save_dir='work_dirs/save/faster/full_det/faster_rcnn_orpn_our_rsp_vit-base-win-rvsa_v3_kvdiff_wsz7_fpn_1x_dota10_lr1e-4_ldr75_dpr15' nproc=1

Evaluating the detection maps predicted by ViTAE-B + RVSA on DIOR-R dataset

CUDA_VISIBLE_DEVICES=0 python tools/test.py configs/obb/oriented_rcnn/vit_base_win/faster_rcnn_orpn_our_rsp_vitae-nc-base-win-rvsa_v3_wsz7_fpn_1x_dior_lr1e-4_ldr75_dpr10.py \
../OBBDetection/work_dirs/faster/faster_rcnn_orpn_our_rsp_vitae-nc-base-win-rvsa_v3_wsz7_fpn_1x_dior_lr1e-4_ldr75_dpr10/latest.pth \
--out work_dirs/save/faster/full_det/faster_rcnn_orpn_our_rsp_vitae-nc-base-win-rvsa_v3_wsz7_fpn_1x_dior_lr1e-4_ldr75_dpr10/det_result.pkl --eval 'mAP' \
--show-dir work_dirs/save/faster/display/faster_rcnn_orpn_our_rsp_vitae-nc-base-win-rvsa_v3_wsz7_fpn_1x_dior_lr1e-4_ldr75_dpr10

Note: the pathes of saved maps and outputs should be constructed before evaluating the DIOR-R testing set.

Training & Evaluation-Segmentation

cd ./Semantic Segmentation

Training and evaluation the UperNet with ViT-B + RVSA backbone on Potsdam dataset:

CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 --master_port=30000 tools/train.py \
configs/vit_base_win/upernet_vit_base_win_rvsa_v3_512x512_160k_potsdam_rgb_dpr10_lr6e5_lrd90_ps16_class5_ignore5.py \
--launcher 'pytorch' --cfg-options 'find_unused_parameters'=True

Note: when training on the LoveDA, please add --no-validate

Inference the LoveDA dataset for online evaluation using the UperNet with ViTAE-B + RVSA $^ \Diamond$ backbone

CUDA_VISIBLE_DEVICES=0 python tools/test.py configs/vit_base_win/upernet_vitae_nc_base_rvsa_v3_kvdiff_wsz7_512x512_160k_loveda_dpr10_lr6e5_lrd90_ps16.py \
../mmsegmentation-master/work_dirs/upernet_vitae_nc_base_rvsa_v3_kvdiff_wsz7_512x512_160k_loveda_dpr10_lr6e5_lrd90_ps16/latest.pth \
--format-only --eval-options imgfile_prefix="work_dirs/display/upernet_vitae_nc_base_rvsa_v3_kvdiff_wsz7_512x512_160k_loveda_dpr10_lr6e5_lrd90_ps16/result" \
--show-dir work_dirs/display/upernet_vitae_nc_base_rvsa_v3_kvdiff_wsz7_512x512_160k_loveda_dpr10_lr6e5_lrd90_ps16/rgb

When finetuning with more than one GPU for detection or segmentation, please use nn.SyncBatchNorm in the NormalCell of ViTAE models.

Citation

If this repo is useful for your research, please consider citation

@ARTICLE{rvsa,
  author={Wang, Di and Zhang, Qiming and Xu, Yufei and Zhang, Jing and Du, Bo and Tao, Dacheng and Zhang, Liangpei},
  journal={IEEE Transactions on Geoscience and Remote Sensing}, 
  title={Advancing Plain Vision Transformer Toward Remote Sensing Foundation Model}, 
  year={2023},
  volume={61},
  number={},
  pages={1-15},
  doi={10.1109/TGRS.2022.3222818}
 }

@ARTICLE{rsp,
  author={Wang, Di and Zhang, Jing and Du, Bo and Xia, Gui-Song and Tao, Dacheng},
  journal={IEEE Transactions on Geoscience and Remote Sensing}, 
  title={An Empirical Study of Remote Sensing Pretraining}, 
  year={2023},
  volume={61},
  number={},
  pages={1-20},
  doi={10.1109/TGRS.2022.3176603}
}

Statement

This project is under MIT licence. For any other questions please contact di.wang at gmail.com .

References

The codes of Pretraining & Scene Classification part mainly from MAE.

Relevant Projects

[1] An Empirical Study of Remote Sensing Pretraining, IEEE TGRS, 2022 | Paper | Github
     Di Wang∗, Jing Zhang∗, Bo Du, Gui-Song Xia and Dacheng Tao