{"id":18417003,"url":"https://github.com/young-geng/m3ae_public","last_synced_at":"2025-04-06T16:15:16.415Z","repository":{"id":48243333,"uuid":"516250002","full_name":"young-geng/m3ae_public","owner":"young-geng","description":"Multimodal Masked Autoencoders (M3AE): A JAX/Flax Implementation","archived":false,"fork":false,"pushed_at":"2025-02-26T19:05:32.000Z","size":29,"stargazers_count":103,"open_issues_count":4,"forks_count":12,"subscribers_count":5,"default_branch":"master","last_synced_at":"2025-03-30T15:10:01.317Z","etag":null,"topics":["computer-vision","flax","jax","natural-language-processing","transformers"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/young-geng.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"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":"2022-07-21T06:22:24.000Z","updated_at":"2025-03-24T14:21:16.000Z","dependencies_parsed_at":"2024-03-23T07:44:27.740Z","dependency_job_id":"939641b8-93ab-4231-95cb-ba18dec11b0b","html_url":"https://github.com/young-geng/m3ae_public","commit_stats":null,"previous_names":[],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/young-geng%2Fm3ae_public","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/young-geng%2Fm3ae_public/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/young-geng%2Fm3ae_public/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/young-geng%2Fm3ae_public/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/young-geng","download_url":"https://codeload.github.com/young-geng/m3ae_public/tar.gz/refs/heads/master","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":247509238,"owners_count":20950232,"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","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","flax","jax","natural-language-processing","transformers"],"created_at":"2024-11-06T04:07:54.948Z","updated_at":"2025-04-06T16:15:16.387Z","avatar_url":"https://github.com/young-geng.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# Multimodal Masked Autoencoders (M3AE): A JAX/Flax Implementation\n\nThis is a JAX/Flax re-implementation for the paper [Multimodal Masked Autoencoders Learn Transferable Representations](https://arxiv.org/abs/2205.14204).\n\n```\n@article{geng2022multimodal,\n  title={Multimodal Masked Autoencoders Learn Transferable Representations},\n  author={Geng, Xinyang and Liu, Hao and Lee, Lisa and Schuurams, Dale and Levine, Sergey and Abbeel, Pieter},\n  journal={arXiv preprint arXiv:2205.14204},\n  year={2022}\n}\n```\n\nThis implementation has been tested on GPU and Google Cloud TPU and supports multi-host training with TPU Pods.\nUnliked the original implementation used for the paper, this implementation also supports the following new\nfeatures:\n* Predicting discretized image tokens from VQGAN as output (similar to BEiT).\n* Training on a combination of paired image-text data (e.g. CC12M) and unpaired text data (e.g. Wikipedia).\n\n## Installation\nIf this is on GPU, replace the following lines in `requirements.txt`\n```\n-f https://storage.googleapis.com/jax-releases/libtpu_releases.html\njax[tpu]==0.3.12\n```\nwith\n```\n--f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\njax[cuda]==0.3.12\n```\nInstall the dependencies with pip and add this repo directory to your `PYTHONPATH` environment variable.\n```\npip install requirements.txt\nexport PYTHONPATH=\"$PYTHONPATH:$(pwd)\"\n```\n\n\n## Running Experiments\nExperiments can be launched via the following commands.\n\nPre-training MAE (image only model) on Conceptual 12M (CC12M)\n```\npython3 -m m3ae.mae_main \\\n    --mae.model_type='large' \\\n    --mae.use_type_embedding=False \\\n    --seed=42 \\\n    --epochs=100 \\\n    --lr_warmup_epochs=5 \\\n    --batch_size=4096 \\\n    --dataloader_n_workers=16 \\\n    --log_freq=500 \\\n    --plot_freq=2000 \\\n    --save_model_freq=10000 \\\n    --lr_peak_value=1.5e-4 \\\n    --weight_decay=0.05 \\\n    --discretized_image=False \\\n    --load_checkpoint='' \\\n    --dataset='cc12m' \\\n    --cc12m_data.path=\"\u003cYOUR DATA HDF5 FILE PATH\u003e\" \\\n    --cc12m_data.image_normalization='cc12m'\n```\n\nPre-training M3AE (image and text model) on Conceptual 12M (CC12M)\n```\npython3 -m m3ae.m3ae_main \\\n    --m3ae.model_type='large' \\\n    --m3ae.image_mask_ratio=0.75 \\\n    --m3ae.text_mask_ratio=0.75 \\\n    --seed=42 \\\n    --epochs=100 \\\n    --lr_warmup_epochs=5 \\\n    --batch_size=4096 \\\n    --discretized_image=False \\\n    --dataloader_n_workers=16 \\\n    --log_freq=500 \\\n    --plot_freq=2000 \\\n    --save_model_freq=10000 \\\n    --image_loss_weight=1.0 \\\n    --text_loss_weight=0.5 \\\n    --lr_peak_value=1.5e-4 \\\n    --weight_decay=0.05 \\\n    --load_checkpoint='' \\\n    --data.path=\"\u003cYOUR DATA HDF5 FILE PATH\u003e\" \\\n    --data.transform_type='pretrain' \\\n    --data.image_normalization='cc12m'\n```\n\nLinear classification on ImageNet for both pre-trained MAE and M3AE\n```\npython3 -m m3ae.linear_main \\\n    --mae.model_type=\"large\" \\\n    --mae.use_type_embedding=True \\\n    --seed=42 \\\n    --epochs=90 \\\n    --batch_size=2048 \\\n    --lr_warmup_epochs=10 \\\n    --discretized_image=False \\\n    --dataloader_n_workers=16 \\\n    --dataloader_shuffle=False \\\n    --log_freq=500 \\\n    --save_model_freq=10000 \\\n    --lr_peak_value=1e-1 \\\n    --weight_decay=0 \\\n    --momentum=0.9 \\\n    --train_data.partition=\"train\" \\\n    --val_data.partition=\"val\" \\\n    --train_data.path=\"\u003cYOUR DATA HDF5 FILE PATH\u003e\" \\\n    --val_data.path=\"\u003cYOUR DATA HDF5 FILE PATH\u003e\" \\\n    --train_data.transform_type=\"linear_prob\" \\\n    --val_data.transform_type=\"test\" \\\n    --load_checkpoint='' \\\n    --load_pretrained=\"\u003cYOUR PRE-TRAINED MODEL PATH\u003e\"\n```\n\nFinetuning on ImageNet for both pre-trained MAE and M3AE\n```\npython3 -m m3ae.finetune_main \\\n    --seed=42 \\\n    --mae.model_type=large \\\n    --mae.drop_path=0.1 \\\n    --weight_decay=0.05 \\\n    --mixup_alpha=0.8 \\\n    --cutmix_alpha=1.0 \\\n    --switch_prob=0.5 \\\n    --label_smoothing=0.1 \\\n    --layer_decay=0.60 \\\n    --clip_gradient=1e9 \\\n    --batch_size=1024 \\\n    --warmup_epochs=5 \\\n    --epochs=100 \\\n    --dataloader_n_workers=16 \\\n    --dataloader_shuffle=False \\\n    --log_freq=500 \\\n    --save_model_freq=10000 \\\n    --lr_peak_value=1e-3 \\\n    --train_data.partition=\"train\" \\\n    --val_data.partition=\"val\" \\\n    --train_data.path=\"\u003cYOUR DATA HDF5 FILE PATH\u003e\" \\\n    --val_data.path=\"\u003cYOUR DATA HDF5 FILE PATH\u003e\" \\\n    --train_data.transform_type=\"finetune\" \\\n    --val_data.transform_type=\"test\" \\\n    --load_pretrained=\"\u003cYOUR PRE-TRAINED MODEL PATH\u003e\"\n```\n\n## HDF5 Data Format\nIn order to facilitate training on cloud, we store all the dataset\nas HDF5 files and read them from cloud storage buckets. For paired image and text\ndataset, the HDF5 data contains two field, `jpg` and `caption`. The `jpg` field\nis an 1D array containing the raw bytes of JPEG encoded images. The `caption`\nfield is an 1D array of utf-8 encoded text. For ImageNet dataset, the image JPEG\nbytes are stored in field `train_jpg` and `val_jpg`, and the integer labels are\nstored in field `train_labels` and `val_labels`. For unpaired text only dataset,\nthe utf-8 encoded text is stored in field `text`.\n\n## Pre-trained Model Weights\n[Pre-trained model weights can be downloaded here](https://www.dropbox.com/scl/fo/hnw0qtdfntl9mxls2xb8g/h?rlkey=vyzg1fm71prznifv6ijjy4t6m\u0026dl=0).\nThe M3AE and MAE models here are trained for 50 epochs on the CC12M dataset using\nthe hyperparameters specified in the paper.\n\nFor converting the pre-trained Jax weights to PyTorch, please refer to [this colab](https://colab.research.google.com/drive/1L6EiUUSkeTzWtB3_K4fR0evn-SDqKWME).\n\n\n## Credits\n* The MAE is heavily inspired by the [original MAE implementation](https://github.com/facebookresearch/mae).\n\n* The VQGAN image tokenizers are from [MaskGiT](https://github.com/google-research/maskgit)\nand [dalle-mini](https://huggingface.co/dalle-mini/vqgan_imagenet_f16_16384).\n\n* The transformer implementation is heavily inspired by [jax-models](https://github.com/DarshanDeshpande/jax-models).\n\n* Some utilities are borrowed from [JaxCQL](https://github.com/young-geng/JaxCQL).\n\n## Contact\nIf you have any questions, please open an issue or contact young.geng@berkeley.edu and hao.liu@berkely.edu.\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fyoung-geng%2Fm3ae_public","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fyoung-geng%2Fm3ae_public","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fyoung-geng%2Fm3ae_public/lists"}