OneScience-Group/Scale-MAE

1---2license: cc-by-nc-4.03language:4- en5tags:6- OneScience7- earth-science8- remote-sensing-representation-learning9- multiscale-remote-sensing10- GSD-scale-modeling11- low-high-frequency-reconstruction12frameworks: PyTorch13datasets:14  - FMoW-RGB15  - RESISC-4516  - EuroSAT17  - UCMerced18  - AID19  - MLRSNet20---21 22<p align="center"><strong><span style="font-size: 30px;">Scale-MAE</span></strong></p>23 24# Model Introduction25 26Scale-MAE is a scale-aware masked autoencoder for multiscale geospatial imagery that learns stable remote sensing image representations through ground-sampling-distance-aware positional encoding, visible-patch encoding, and low- and high-frequency target reconstruction.27 28Paper: Scale-MAE: A Scale-Aware Masked Autoencoder for Multiscale Geospatial Representation Learning  29https://arxiv.org/abs/2212.1453230 31# Model Description32 33Scale-MAE was proposed by research teams at NASA's Jet Propulsion Laboratory and Stanford University. The model is trained using multiscale geospatial imagery such as FMoW-RGB. The model is suitable for tasks such as remote sensing image representation learning, scene classification, and building segmentation.34 35# Applicable Scenarios36 37| Scenario | Description |38| :---: | :--- |39| Multiscale remote sensing pre-training | Use paired low-resolution and high-resolution `BCHW` imagery with GSD metadata. |40| Scene classification | Perform kNN transfer evaluation through reusable CLS features. |41| Building segmentation | Transfer scale-aware representations to building semantic-segmentation tasks such as SpaceNet and fine-tune them. |42| Low- and high-frequency reconstruction | Use area resampling and band-pass targets to evaluate scale sensitivity. |43| Local quick validation | Use synthetic data to check data loading, training, inference, and evaluation. |44| Multi-GPU training | Launch distributed data-parallel training through `torchrun`. |45 46# Usage Instructions47 48## 1. OneCode Usage49 50Experience intelligent one-click AI4S programming through the OneCode online environment:51 52[Experience intelligent one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)53 54## 2. Manual Installation and Usage55 56**Hardware Requirements**57 58- GPU or DCU execution is recommended.59- CPU can be used to validate the workflow with the current default small configuration.60- DCU users need to install a DTK version matching the cluster in advance. DTK 25.04.2 or later is recommended.61 62### Download the Model Package63 64```bash65hf download OneScience-Group/Scale-MAE --local-dir ./Scale-MAE66cd Scale-MAE67```68 69### Install the Runtime Environment70 71**DCU Environment**72 73```bash74conda create -n onescience311 python=3.11 -y75conda activate onescience31176pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai77```78 79**GPU Environment**80 81```bash82conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=1283conda activate onescience31184pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai85```86 87### Training Data Introduction88 89The paper uses multiscale geospatial imagery such as FMoW-RGB for pre-training and evaluates on tasks including RESISC-45, UCMerced, EuroSAT, AID, MLRSNet, and SpaceNet. Data files contain `images` `[B,C,H,W]`, `targets` `[B,C,Ht,Wt]`, `gsd` `[B]`, and `labels` `[B]`.90 91Synthetic data is used by default:92 93```bash94python scripts/fake_data.py95```96 97When using real data, do not run `fake_data.py`. First organize the data into the following directories and fields, and replace the files under `data/` generated by the synthetic-data script:98 99```text100data/train.npz101data/test.npz102```103 104Each NPZ file contains at least:105 106```text107images:  float32 [N,C,input_size,input_size]108targets: float32 [N,C,target_size,target_size]109gsd:     float32 [N]110labels:  int64   [N]111```112 113Here, `images` is the model input, `targets` is the target-resolution imagery corresponding to the input scene, `gsd` is the ground sampling distance for each sample in meters per pixel, and `labels` is used for kNN feature evaluation. The channel count, input size, target size, and GSD range of real data must be consistent with `conf/config.yaml` and the model configuration; cropping, registration, channel organization, and numerical normalization should be completed before generating the NPZ files.114 115Modify `input_size`, `target_size`, `channels`, `gsd_values`, and paths in `conf/config.yaml` according to the real data. After completing data preparation, continue to use the unified training, inference, and evaluation commands below; if other file locations are needed, override the default paths through script arguments.116 117### Training118 119Single GPU:120 121```bash122python scripts/train.py123```124 125Multiple GPUs:126 127```bash128torchrun --nproc_per_node=8 scripts/train.py129```130 131Training outputs:132 133```text134result/checkpoints/scalemae.pt135result/training/metrics.json136```137 138Training outputs include a model checkpoint that can be used for subsequent inference and feature extraction, as well as training metrics reflecting changes in overall, low-frequency, and high-frequency reconstruction losses, facilitating training-state preservation and analysis of model optimization.139 140AdamW uses betas `(0.9, 0.95)` and includes gradient accumulation, AMP, warmup, and cosine decay.141 142### Trained Weights143 144This repository provides weights trained on multiscale geospatial imagery in the `weight/` folder. The weight files will be uploaded soon.145 146### Inference147 148```bash149python scripts/inference.py150```151 152Inference results are output to:153 154```text155result/output/reconstruction.npz156```157 158### Evaluation and Visualization159 160```bash161python scripts/result.py162```163 164Evaluation and visualization outputs are saved to:165 166```text167result/evaluation/metrics.json168result/evaluation/features.npy169result/evaluation/bandpass_reconstruction.png170result/evaluation/frequency_error.png171result/evaluation/gsd_reconstruction_error.png172result/evaluation/gsd_knn_accuracy.png173```174 175Evaluation results comprehensively reflect overall and low- and high-frequency reconstruction quality, scale adaptability under different GSD values, and the kNN classification capability of representation features, while reconstruction comparisons and scale-variation curves demonstrate the model's ability to process multiscale geospatial imagery. The current results are based on a small amount of synthetic data and are mainly used to confirm that the training, inference, evaluation, and visualization workflows operate normally.176 177# OneScience Official Information178 179| Platform | OneScience Main Repository | Skills Repository |180| --- | --- | --- |181| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |182| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |183 184# Citation and License185 186This repository is a reproduction of the original Scale-MAE paper.187