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 |

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.

3. 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