{"id":31055391,"url":"https://github.com/tonylianlong/crossmae","last_synced_at":"2025-09-15T04:44:19.721Z","repository":{"id":219164536,"uuid":"747424402","full_name":"TonyLianLong/CrossMAE","owner":"TonyLianLong","description":"Official Implementation of the CrossMAE paper: Rethinking Patch Dependence for Masked Autoencoders","archived":false,"fork":false,"pushed_at":"2025-04-10T05:56:37.000Z","size":1574,"stargazers_count":107,"open_issues_count":1,"forks_count":6,"subscribers_count":3,"default_branch":"main","last_synced_at":"2025-04-10T06:38:42.126Z","etag":null,"topics":["computer-vision","deep-learning","mae","masked-autoencoder","self-supervised-learning"],"latest_commit_sha":null,"homepage":"https://crossmae.github.io/","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"other","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/TonyLianLong.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":"CODE_OF_CONDUCT.md","threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2024-01-23T22:34:59.000Z","updated_at":"2025-04-10T05:56:41.000Z","dependencies_parsed_at":"2024-01-28T01:04:50.186Z","dependency_job_id":"e6a747c5-376b-48c8-a94a-e2620e25545c","html_url":"https://github.com/TonyLianLong/CrossMAE","commit_stats":null,"previous_names":["tonylianlong/crossmae"],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/TonyLianLong/CrossMAE","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/TonyLianLong%2FCrossMAE","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/TonyLianLong%2FCrossMAE/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/TonyLianLong%2FCrossMAE/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/TonyLianLong%2FCrossMAE/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/TonyLianLong","download_url":"https://codeload.github.com/TonyLianLong/CrossMAE/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/TonyLianLong%2FCrossMAE/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":275207721,"owners_count":25423895,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","status":"online","status_checked_at":"2025-09-15T02:00:09.272Z","response_time":75,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["computer-vision","deep-learning","mae","masked-autoencoder","self-supervised-learning"],"created_at":"2025-09-15T04:44:04.951Z","updated_at":"2025-09-15T04:44:19.707Z","avatar_url":"https://github.com/TonyLianLong.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"## CrossMAE: Rethinking Patch Dependence for Masked Autoencoders\nby \u003ca href=\"https://max-fu.github.io\"\u003eLetian Fu*\u003c/a\u003e, \u003ca href=\"https://tonylian.com\"\u003eLong Lian*\u003c/a\u003e, \u003ca href=\"https://renwang435.github.io\"\u003eRenhao Wang\u003c/a\u003e, \u003ca href=\"https://bfshi.github.io\"\u003eBaifeng Shi\u003c/a\u003e, \u003ca href=\"https://people.eecs.berkeley.edu/~xdwang\"\u003eXudong Wang\u003c/a\u003e, \u003ca href=\"https://www.adamyala.org\"\u003eAdam Yala†\u003c/a\u003e, \u003ca href=\"https://people.eecs.berkeley.edu/~trevor\"\u003eTrevor Darrell†\u003c/a\u003e, \u003ca href=\"https://people.eecs.berkeley.edu/~efros\"\u003eAlexei A. Efros†\u003c/a\u003e, \u003ca href=\"https://goldberg.berkeley.edu\"\u003eKen Goldberg†\u003c/a\u003e at UC Berkeley and UCSF\n\n[[Paper](https://openreview.net/forum?id=JT2KMuo2BV)] | [[Project Page](https://crossmae.github.io/)] | [[Citation](#citation)]\n\n\n\u003cp align=\"center\"\u003e\n  \u003cimg src=\"https://crossmae.github.io/crossmae2.jpg\" width=\"800\"\u003e\n\u003c/p\u003e\n\nThis is a PyTorch implementation of the CrossMAE paper [Rethinking Patch Dependence for Masked Autoencoders](https://crossmae.github.io/). The code is based on the original [MAE](https://github.com/facebookresearch/mae) repo. The codebase supports CrossMAE and MAE, with `timm==0.9.7`, `torch==2.0.0`, and flash-attn 2.\n\n## Models\nThe encoder part of CrossMAE matches exactly with MAE. Therefore, we use the same code for fine-tuning. We also encourage you to try CrossMAE checkpoints in your downstream applications. These models are trained on ImageNet-1k for 800 epochs (except that 448 models are trained for 400 epochs), with masking ratio and kept mask ratio both set to 0.75, except that ViT-H is with masking ratio 0.75 and kept mask ratio 0.25.\n\n\u003ctable\u003e\u003ctbody\u003e\n\u003c!-- START TABLE --\u003e\n\u003c!-- TABLE HEADER --\u003e\n\u003cth valign=\"bottom\"\u003e\u003c/th\u003e\n\u003cth valign=\"bottom\"\u003eViT-Small\u003c/th\u003e\n\u003cth valign=\"bottom\"\u003eViT-Base\u003c/th\u003e\n\u003cth valign=\"bottom\"\u003eViT-Base\u003csub\u003e448\u003c/sub\u003e\u003c/th\u003e\n\u003cth valign=\"bottom\"\u003eViT-Large\u003c/th\u003e\n\u003cth valign=\"bottom\"\u003eViT-Huge\u003c/th\u003e\n\u003c!-- TABLE BODY --\u003e\n\u003ctr\u003e\u003ctd align=\"left\"\u003epretrained checkpoint\u003c/td\u003e\n\u003ctd align=\"center\"\u003e\u003ca href='https://huggingface.co/longlian/CrossMAE/resolve/main/vits-mr0.75-kmr0.75-dd12/imagenet-mae-cross-vits-pretrain-wfm-mr0.75-kmr0.75-dd12-ep800-ui.pth?download=true'\u003edownload\u003c/a\u003e\u003c/td\u003e\n\u003ctd align=\"center\"\u003e\u003ca href='https://huggingface.co/longlian/CrossMAE/resolve/main/vitb-mr0.75-kmr0.75-dd12/imagenet-mae-cross-vitb-pretrain-wfm-mr0.75-kmr0.75-dd12-ep800-ui.pth?download=true'\u003edownload\u003c/a\u003e\u003c/td\u003e\n\u003ctd align=\"center\"\u003e\u003ca href='https://huggingface.co/longlian/CrossMAE/resolve/main/vitb-mr0.75-kmr0.75-dd12-448-400/imagenet-mae-cross-vitb-pretrain-wfm-mr0.75-kmr0.25-dd12-ep400-ui-res-448.pth?download=true'\u003edownload\u003c/a\u003e\u003c/td\u003e\n\u003ctd align=\"center\"\u003e\u003ca href='https://huggingface.co/longlian/CrossMAE/resolve/main/vitl-mr0.75-kmr0.75-dd12/imagenet-mae-cross-vitl-pretrain-wfm-mr0.75-kmr0.75-dd12-ep800-ui.pth?download=true'\u003edownload\u003c/a\u003e\u003c/td\u003e\n\u003ctd align=\"center\"\u003e\u003ca href='https://huggingface.co/longlian/CrossMAE/resolve/main/vith-mr0.75-kmr0.25-dd12/imagenet-mae-cross-vith-pretrain-wfm-mr0.75-kmr0.25-dd12-ep800-ui.pth?download=true'\u003edownload\u003c/a\u003e\u003c/td\u003e\n\u003c/tr\u003e\n\u003ctr\u003e\u003ctd align=\"left\"\u003efine-tuned checkpoint\u003c/td\u003e\n\u003ctd align=\"center\"\u003e\u003ca href='https://huggingface.co/longlian/CrossMAE/resolve/main/vits-mr0.75-kmr0.75-dd12/imagenet-mae-cross-vits-finetune-wfm-mr0.75-kmr0.75-dd12-ep800-ui.pth?download=true'\u003edownload\u003c/a\u003e\u003c/td\u003e\n\u003ctd align=\"center\"\u003e\u003ca href='https://huggingface.co/longlian/CrossMAE/resolve/main/vitb-mr0.75-kmr0.75-dd12/imagenet-mae-cross-vitb-finetune-wfm-mr0.75-kmr0.75-dd12-ep800-ui.pth?download=true'\u003edownload\u003c/a\u003e\u003c/td\u003e\n\u003ctd align=\"center\"\u003e\u003ca href='https://huggingface.co/longlian/CrossMAE/resolve/main/vitb-mr0.75-kmr0.75-dd12-448-400/imagenet-mae-cross-vitb-finetune-wfm-mr0.75-kmr0.25-dd12-ep400-ui-res-448.pth?download=true'\u003edownload\u003c/a\u003e\u003c/td\u003e\n\u003ctd align=\"center\"\u003e\u003ca href='https://huggingface.co/longlian/CrossMAE/resolve/main/vitl-mr0.75-kmr0.75-dd12/imagenet-mae-cross-vitl-finetune-wfm-mr0.75-kmr0.75-dd12-ep800-ui.pth?download=true'\u003edownload\u003c/a\u003e\u003c/td\u003e\n\u003ctd align=\"center\"\u003e\u003ca href='https://huggingface.co/longlian/CrossMAE/resolve/main/vith-mr0.75-kmr0.25-dd12/imagenet-mae-cross-vith-finetune-wfm-mr0.75-kmr0.25-dd12-ep800-ui.pth?download=true'\u003edownload\u003c/a\u003e\u003c/td\u003e\n\u003c/tr\u003e\n\u003ctr\u003e\u003ctd align=\"left\"\u003e\u003cb\u003eReference ImageNet accuracy (ours)\u003c/b\u003e\u003c/td\u003e\n\u003ctd align=\"center\"\u003e\u003cb\u003e79.318\u003c/b\u003e\u003c/td\u003e\n\u003ctd align=\"center\"\u003e\u003cb\u003e83.722\u003c/b\u003e\u003c/td\u003e\n\u003ctd align=\"center\"\u003e\u003cb\u003e84.598\u003c/b\u003e\u003c/td\u003e\n\u003ctd align=\"center\"\u003e\u003cb\u003e85.432\u003c/b\u003e\u003c/td\u003e\n\u003ctd align=\"center\"\u003e\u003cb\u003e86.256\u003c/b\u003e\u003c/td\u003e\n\u003c/tr\u003e\n\u003ctr\u003e\u003ctd align=\"left\"\u003eMAE ImageNet accuracy (baseline)\u003c/td\u003e\n\u003ctd align=\"center\"\u003e\u003c/td\u003e\n\u003ctd align=\"center\"\u003e\u003c/td\u003e\n\u003ctd align=\"center\"\u003e84.8\u003c/td\u003e\n\u003ctd align=\"center\"\u003e\u003c/td\u003e\n\u003ctd align=\"center\"\u003e85.9\u003c/td\u003e\n\u003c/tr\u003e\n\u003c/tbody\u003e\u003c/table\u003e\n\n## Train CrossMAE on **one single RTX 4090**\nWith the efficiency of CrossMAE, it's possible to train CrossMAE on **one single RTX 4090** on a personal computer. The CPU is i9-14900k, with 96GB RAM.\n\n\u003cdetails\u003e\n  \u003csummary\u003eInstructions and trained models\u003c/summary\u003e\n\nThe training and fine-tuning command (with `${IMAGENET_DIR}` the directory for imagenet, ViT-S as an example):\n```sh\nCUDA_VISIBLE_DEVICES=0 OMP_NUM_THREADS=1 torchrun --nproc_per_node=1 --master_port 2780 main_pretrain.py --batch_size 512 --accum_iter 8 --model mae_vit_small_patch16 --norm_pix_loss --blr 1.5e-4 --weight_decay 0.05 --data_path ${IMAGENET_DIR} --num_workers 16 --multi_epochs_dataloader --output_dir output/imagenet-crossmae-vits-pretrain-wfm-mr0.75-kmr0.25-dd12-ep800 --cross_mae --weight_fm --decoder_depth 12 --mask_ratio 0.75 --kept_mask_ratio 0.75 --epochs 800 --warmup_epochs 40 --use_input\n\nCUDA_VISIBLE_DEVICES=0 OMP_NUM_THREADS=1 torchrun --nproc_per_node=1 --master_port 2860 main_finetune.py --batch_size 512 --accum_iter 2 --model vit_small_patch16 --finetune output/imagenet-crossmae-vits-pretrain-wfm-mr0.75-kmr0.25-dd12-ep800/checkpoint.pth --epoch 100 --blr 5e-4 --layer_decay 0.65 --weight_decay 0.05 --drop_path 0.1 --reprob 0.25 --mixup 0.8 --cutmix 1.0 --dist_eval --data_path ${IMAGENET_DIR} --num_workers 12 --output_dir output/imagenet-crossmae-vits-finetune-wfm-mr0.75-kmr0.25-dd12-ep800 --multi_epochs_dataloader\n# Reference results:\n# * Acc@1 79.462 Acc@5 94.864 loss 0.907\n```\n\n\u003ctable\u003e\n\u003ctbody\u003e\n  \u003c!-- START TABLE --\u003e\n  \u003c!-- TABLE HEADER --\u003e\n  \u003ctr valign=\"bottom\"\u003e\n    \u003cth align=\"left\"\u003epretrained checkpoint\u003c/th\u003e\n    \u003cth align=\"left\"\u003efine-tuned checkpoint\u003c/th\u003e\n    \u003cth align=\"left\"\u003ereference ImageNet accuracy\u003c/th\u003e\n  \u003c/tr\u003e\n  \u003c!-- TABLE BODY --\u003e\n  \u003ctr\u003e\n    \u003ctd align=\"center\"\u003e\u003ca href='https://huggingface.co/longlian/CrossMAE/resolve/main/vits-mr0.75-kmr0.25-dd12/imagenet-mae-cross-vits-pretrain-wfm-mr0.75-kmr0.75-dd12-ep800-ui.pth?download=true'\u003edownload\u003c/a\u003e\u003c/td\u003e\n    \u003ctd align=\"center\"\u003e\u003ca href='https://huggingface.co/longlian/CrossMAE/resolve/main/vits-mr0.75-kmr0.25-dd12/imagenet-mae-cross-vits-finetune-wfm-mr0.75-kmr0.75-dd12-ep800-ui.pth?download=true'\u003edownload\u003c/a\u003e\u003c/td\u003e\n    \u003ctd align=\"center\"\u003e79.462\u003c/td\u003e\n  \u003c/tr\u003e\n\u003c/tbody\u003e\n\u003c/table\u003e\n\n\u003c/details\u003e\n\n## Instructions\nPlease install the dependencies in `requirements.txt`:\n```sh\n# Optionally create a conda environment\nconda create -n crossmae python=3.10 -y\nconda activate crossmae\n# Install dependencies\npip install -r requirements.txt\n```\n\n### Pre-training CrossMAE\nTo pre-train ViT-Base, run the following on 4 GPUs:\n```sh\nCUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nproc_per_node=4 --master_port 1234 main_pretrain.py --batch_size 1024 --model mae_vit_base_patch16 --norm_pix_loss --blr 1.5e-4 --weight_decay 0.05 --data_path ${IMAGENET_DIR} --num_workers 20 --enable_flash_attention2 --multi_epochs_dataloader --output_dir output/imagenet-crossmae-vitb-pretrain-wfm-mr0.75-kmr0.25-dd12-ep800 --cross_mae --weight_fm --decoder_depth 12 --mask_ratio 0.75 --kept_mask_ratio 0.25 --epochs 800 --warmup_epochs 40 --use_input\n```\n\nTo train ViT-Small or ViT-Large, set `--model mae_vit_small_patch16` or `--model mae_vit_large_patch16`. You can use `--accum_iter` to perform gradient accumulation if your hardware could not fit the batch size. [FlashAttention 2](https://github.com/Dao-AILab/flash-attention) should be installed with `pip install flash-attn --no-build-isolation`.\n\n### Fine-tuning CrossMAE\nTo pre-train ViT-Base, run the following on 4 GPUs:\n```sh\nCUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nproc_per_node=4 --master_port 1234 main_finetune.py --batch_size 256 --model vit_base_patch16 --finetune output/imagenet-crossmae-vitb-pretrain-wfm-mr0.75-kmr0.25-dd12-ep800/checkpoint.pth --epoch 100 --blr 5e-4 --layer_decay 0.65 --weight_decay 0.05 --drop_path 0.1 --reprob 0.25 --mixup 0.8 --cutmix 1.0 --dist_eval --data_path ${IMAGENET_DIR} --output_dir output/imagenet-crossmae-vitb-finetune-wfm-mr0.75-kmr0.25-dd12-ep800 --enable_flash_attention2 --multi_epochs_dataloader\n```\n\n## Evaluation\nEvaluate ViT-Base in a single GPU (`${IMAGENET_DIR}` is a directory containing `{train, val}` sets of ImageNet). `${FINETUNED_CHECKPOINT_PATH}` is the path to the fine-tuned checkpoint:\n```sh\npython main_finetune.py --eval --resume ${FINETUNED_CHECKPOINT_PATH} --model vit_base_patch16 --batch_size 16 --data_path ${IMAGENET_DIR}\n```\nThis should give:\n```\n* Acc@1 83.722 Acc@5 96.686 loss 0.729\n```\n\nYou could replace `vit_base_patch16` with `vit_small_patch16` or `vit_large_patch16` to evaluate ViT-S or ViT-L. To work with 448 input resolution, please append `--input_size 448` to the command line.\n\n### License\n\nThis project is under the CC-BY-NC 4.0 license. See [LICENSE](LICENSE) for details.\n\n## Citation\nPlease give us a star 🌟 on Github to support us!\n\nPlease cite our work if you find our work inspiring or use our code in your work:\n```\n@article{\n    fu2025rethinking,\n    title={Rethinking Patch Dependence for Masked Autoencoders},\n    author={Letian Fu and Long Lian and Renhao Wang and Baifeng Shi and XuDong Wang and Adam Yala and Trevor Darrell and Alexei A Efros and Ken Goldberg},\n    journal={Transactions on Machine Learning Research},\n    issn={2835-8856},\n    year={2025},\n    url={https://openreview.net/forum?id=JT2KMuo2BV},\n    note={}\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Ftonylianlong%2Fcrossmae","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Ftonylianlong%2Fcrossmae","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Ftonylianlong%2Fcrossmae/lists"}