OneScience-Group/SatMAE

1---2license: cc-by-nc-4.03language:4- en5tags:6- OneScience7- earth-science8- remote-sensing-representation-learning9- satellite-imagery10- temporal-remote-sensing11- multispectral-remote-sensing12frameworks: PyTorch13datasets: []14---15 16<p align="center">17  <strong>18    <span style="font-size: 30px;">SatMAE</span>19  </strong>20</p>21 22# Model Introduction23 24SatMAE is a masked autoencoder pre-training model for temporal and multispectral satellite imagery. It learns remote sensing image representations through temporal positional encoding, spectral group encoding, and independent masking across temporal or spectral dimensions, and is mainly used to improve performance on tasks such as satellite image classification, land-cover classification, multi-label classification, and semantic segmentation when labeled data is limited.25 26Paper: SatMAE: Pre-training Transformers for Temporal and Multi-Spectral Satellite Imagery  27https://arxiv.org/abs/2207.0805128 29# Model Description30 31SatMAE was proposed by a research team at Stanford University. The model is pre-trained using fMoW satellite imagery, temporal satellite imagery, and Sentinel-2 multispectral imagery. The model is suitable for remote sensing tasks such as satellite image classification, land-cover classification, multi-label classification, and semantic segmentation.32 33# Applicable Scenarios34 35| Scenario | Description |36| :---: | :--- |37| Temporal satellite image pre-training | Train SatMAE using timestamped `BTCHW` data. |38| Multispectral satellite image pre-training | Train SatMAE using `BCHW` data and a spectral grouping configuration. |39| Local quick validation | Use synthetic data to check data loading, training, inference, and evaluation. |40| ModelScope/OneCode operation | Run scripts after downloading it as a standalone model package. |41| Multi-GPU training | Launch multi-process training through `torchrun`. |42 43# Usage Instructions44 45## 1. OneCode Usage46 47Experience intelligent one-click AI4S programming through the OneCode online environment:48 49[Experience intelligent one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)50 51## 2. Manual Installation and Usage52 53**Hardware Requirements**54 55- GPU or DCU execution is recommended.56- CPU can be used for imports and small-configuration connectivity validation, but full training and inference are slow.57- DCU users need to install DTK in advance. DTK 25.04.2 or later, or the OneScience-recommended version matching the current cluster, is recommended.58 59### Download the Model Package60 61```bash62hf download OneScience-Group/SatMAE --local-dir ./SatMAE63cd SatMAE64```65 66### Install the Runtime Environment67 68**DCU Environment**69 70```bash71# Activate DTK and CONDA first72conda create -n onescience311 python=3.11 -y73conda activate onescience31174# Supports installation with uv75pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai76```77 78**GPU Environment**79 80```bash81# Activate CONDA first82conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=1283conda activate onescience31184# Supports installation with uv85pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai86```87 88### Training Data Introduction89 90The temporal experiment in the paper uses fMoW RGB, with input consisting of RGB image sequences of length 3 from the same location; the multispectral experiment uses fMoW-Sentinel. In temporal-mode NPZ files, `images` is `[B,T,C,H,W]`, `timestamps` is `[B,T,3]`, with the three fields being `year - 2002`, `month - 1`, and `hour` in order, and `labels` is `[B]`. The model also accepts continuous scalar time in `[B,T]` format.91 92Small synthetic data following the same protocol is used by default:93 94```bash95python scripts/fake_data.py96```97 98Run the command above when using synthetic data. When using real data, do not run `fake_data.py`; save the data as `data/train.npz` and `data/test.npz`, and modify `conf/config.yaml` according to the actual protocol.99 100```text101images:     float32 [N,T,C,H,W]102timestamps: float32 [N,T,3]103labels:     int64   [N]104```105 106The three fields of `timestamps` are `year - 2002`, `month - 1`, and `hour` in order; continuous scalar time can also use the `[N,T]` format. The data should have completed size processing, channel ordering, temporal sorting, and numerical normalization.107 108### Training109 110Single GPU:111 112```bash113python scripts/train.py114```115 116Multiple GPUs:117 118```bash119torchrun --nproc_per_node=8 scripts/train.py120```121 122Training outputs:123 124```text125result/checkpoints/satmae.pt126result/training/metrics.json127```128 129The default configuration uses small synthetic data to validate the training workflow. Formal training should use the Base, Large, or Huge paper architecture presets provided by the code, together with real data and the paper's training budget.130 131### Trained Weights132 133This repository provides weights trained on fMoW RGB temporal satellite imagery and fMoW-Sentinel multispectral satellite imagery in the `weight/` folder. The weight files will be uploaded soon and are expected to be completed in the near future.134 135### Inference136 137```bash138python scripts/inference.py139```140 141Inference results are output to:142 143```text144result/output/reconstruction.npz145```146 147### Evaluation and Visualization148 149```bash150python scripts/result.py151```152 153Evaluation and visualization outputs are saved to:154 155```text156result/evaluation/metrics.json157result/evaluation/temporal_frame_reconstruction.png158result/evaluation/temporal_reconstruction_error.png159result/evaluation/spectral_band_reconstruction.png160```161 162Evaluation results include overall and masked-patch MSE, overall and masked-patch MSE for each temporal frame, and reconstruction MSE for each input channel. Synthetic-data results are only used to validate the engineering workflow and do not represent downstream transfer performance from the paper.163 164# OneScience Official Information165 166| Platform | OneScience Main Repository | Skills Repository |167| --- | --- | --- |168| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |169| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |170 171# Citation and License172 173This repository is a reproduction of the original SatMAE paper.174