Updates | Introduction | Results & Models | Usage |
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;
Our models have been supported by LuoJiaNET, please refer to RS-Vision-Foundation-Models for more details.
The early access is available! TGRS link
The arXiv has been updated! arXiv link
The paper has been accepted by IEEE TGRS!
The codes, configs and training logs of segmentation in fintuning are released!
The codes, configs and training logs of detection in fintuning are released!
The codes of pretraining and classification in fintuning are released!
The codes and training logs of the VSA have been released, which is the foundation of our RVSA.
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.
Pretrain | Backbone | Input size | Params (M) | Pretrained model |
---|---|---|---|---|
MAE | ViT-B | 224 × 224 | 86 | Weights |
MAE | ViTAE-B | 224 × 224 | 89 | Weights |
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 |
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 |
1x | 78.99 | Config | Log | Model |
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 |
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 |
1x | 81.18 | Config | Log | Model |
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 |
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 |
1x | 71.05 | Config | Log | Model |
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 |
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 |
99.50 | 97.01 | 98.50 | 93.92 | 95.66 |
Model | Model | Model | Model | Model |
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 |
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 |
512 × 512 | 160k | 91.15 | Config | Log | Model |
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 |
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 |
896 × 896 | 160k | 64.49 | Config | Log | Model |
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 |
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 |
512 × 512 | 160k | 52.44 | Config | Log | Model |
Environment:
- Python 3.8.5
- Pytorch 1.9.0+cu111
- torchvision 0.10.0+cu111
- timm 0.4.12
- mmcv-full 1.4.1
-
Preparing the MillionAID: Download the MillionAID. Here, we use previous
train_labels.txt
andvalid_labels.txt
of the ViTAE-Transformer-Remote-Sensing, 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 codesMAEPretrain_SceneClassification/util/datasets.py/class MillionAIDDataset
. -
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.
- 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'
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'
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 ViTAE-Transformer-Remote-Sensing
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
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
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
Predicting the saving detection map using ViT-B + RVSA
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.
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
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.
If this repo is useful for your research, please consider citation
@ARTICLE{wang_vitrvsa_2022,
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 Towards Remote Sensing Foundation Model},
year={2022},
volume={},
number={},
pages={1-1},
doi={10.1109/TGRS.2022.3222818}
}
@ARTICLE{wang_rsp_2022,
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={2022},
volume={},
number={},
pages={1-1},
doi={10.1109/TGRS.2022.3176603}}
This project is under MIT licence. For any other questions please contact di.wang at gmail.com .
The codes of Pretraining & Scene Classification part mainly from MAE.
[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