m3ae_public - GitLite

KO | EN

gitlite — search

#javascript #react #ai #llm #golang #android #deep-learning #nodejs #java #windows #awesome #vue #pytorch #openai

m3ae_public

(★ 111)

Multimodal Masked Autoencoders (M3AE): A JAX/Flax Implementation

파일 탐색기

최종 버전 다운로드 (.zip)

- __init__.py
- data.py
- finetune_main.py
- jax_utils.py
- linear_main.py
- m3ae_main.py
- mae_main.py
- model.py
- utils.py
- vqgan.py
.gitignore
LICENSE
README.md
requirements.txt

# 설치 가이드

1. 코드 내려받기

git clone https://github.com/young-geng/m3ae_public

깃허브에서 프로젝트 코드 전체를 내 컴퓨터로 내려받습니다.

cd m3ae_public

방금 내려받은 프로젝트 폴더 안으로 이동합니다.

2. 공식 설치 스크립트

쉬움 추천

사전 준비물

Python 3 pip 명령어를 쓰려면 Python이 필요합니다.

pip install requirements.txt

PyPI에 배포된 패키지를 바로 설치합니다. 소스 클론이 필요 없습니다.

✅ 설치 후 새 터미널을 열고, 프로그램의 버전 확인 명령(예: --version)으로 정상 설치됐는지 확인하세요.

이 레포의 README에 적힌 실제 명령어를 그대로 가져왔습니다.

3. Python

쉬움

사전 준비물

Git GitHub에서 프로젝트 코드를 내려받으려면 필요합니다.
Python 3 설치 시 'Add Python to PATH' 옵션을 꼭 체크하세요 (Windows).

pip install requirements.txt

PyPI에 배포된 패키지를 바로 설치합니다. 소스 클론이 필요 없습니다.

python3 -m m3ae.mae_main \

파이썬 스크립트(또는 모듈)를 실행합니다.

python3 -m m3ae.m3ae_main \

파이썬 스크립트(또는 모듈)를 실행합니다.

python3 -m m3ae.linear_main \

파이썬 스크립트(또는 모듈)를 실행합니다.

python3 -m m3ae.finetune_main \

파이썬 스크립트(또는 모듈)를 실행합니다.

✅ 에러 메시지 없이 실행되고 터미널에 안내 문구가 출력되면 정상입니다.

이 레포의 README에 적힌 실제 명령어를 그대로 가져왔습니다.

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# Multimodal Masked Autoencoders (M3AE): A JAX/Flax Implementation

This is a JAX/Flax re-implementation for the paper [Multimodal Masked Autoencoders Learn Transferable Representations](https://arxiv.org/abs/2205.14204).

```
@article{geng2022multimodal,
  title={Multimodal Masked Autoencoders Learn Transferable Representations},
  author={Geng, Xinyang and Liu, Hao and Lee, Lisa and Schuurams, Dale and Levine, Sergey and Abbeel, Pieter},
  journal={arXiv preprint arXiv:2205.14204},
  year={2022}
}
```

This implementation has been tested on GPU and Google Cloud TPU and supports multi-host training with TPU Pods.
Unliked the original implementation used for the paper, this implementation also supports the following new
features:
* Predicting discretized image tokens from VQGAN as output (similar to BEiT).
* Training on a combination of paired image-text data (e.g. CC12M) and unpaired text data (e.g. Wikipedia).

## Installation
If this is on GPU, replace the following lines in `requirements.txt`
```
-f https://storage.googleapis.com/jax-releases/libtpu_releases.html
jax[tpu]==0.3.12
```
with
```
--f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
jax[cuda]==0.3.12
```
Install the dependencies with pip and add this repo directory to your `PYTHONPATH` environment variable.
```
pip install requirements.txt
export PYTHONPATH="$PYTHONPATH:$(pwd)"
```

## Running Experiments
Experiments can be launched via the following commands.

Pre-training MAE (image only model) on Conceptual 12M (CC12M)
```
python3 -m m3ae.mae_main \
    --mae.model_type='large' \
    --mae.use_type_embedding=False \
    --seed=42 \
    --epochs=100 \
    --lr_warmup_epochs=5 \
    --batch_size=4096 \
    --dataloader_n_workers=16 \
    --log_freq=500 \
    --plot_freq=2000 \
    --save_model_freq=10000 \
    --lr_peak_value=1.5e-4 \
    --weight_decay=0.05 \
    --discretized_image=False \
    --load_checkpoint='' \
    --dataset='cc12m' \
    --cc12m_data.path="<YOUR DATA HDF5 FILE PATH>" \
    --cc12m_data.image_normalization='cc12m'
```

Pre-training M3AE (image and text model) on Conceptual 12M (CC12M)
```
python3 -m m3ae.m3ae_main \
    --m3ae.model_type='large' \
    --m3ae.image_mask_ratio=0.75 \
    --m3ae.text_mask_ratio=0.75 \
    --seed=42 \
    --epochs=100 \
    --lr_warmup_epochs=5 \
    --batch_size=4096 \
    --discretized_image=False \
    --dataloader_n_workers=16 \
    --log_freq=500 \
    --plot_freq=2000 \
    --save_model_freq=10000 \
    --image_loss_weight=1.0 \
    --text_loss_weight=0.5 \
    --lr_peak_value=1.5e-4 \
    --weight_decay=0.05 \
    --load_checkpoint='' \
    --data.path="<YOUR DATA HDF5 FILE PATH>" \
    --data.transform_type='pretrain' \
    --data.image_normalization='cc12m'
```

Linear classification on ImageNet for both pre-trained MAE and M3AE
```
python3 -m m3ae.linear_main \
    --mae.model_type="large" \
    --mae.use_type_embedding=True \
    --seed=42 \
    --epochs=90 \
    --batch_size=2048 \
    --lr_warmup_epochs=10 \
    --discretized_image=False \
    --dataloader_n_workers=16 \
    --dataloader_shuffle=False \
    --log_freq=500 \
    --save_model_freq=10000 \
    --lr_peak_value=1e-1 \
    --weight_decay=0 \
    --momentum=0.9 \
    --train_data.partition="train" \
    --val_data.partition="val" \
    --train_data.path="<YOUR DATA HDF5 FILE PATH>" \
    --val_data.path="<YOUR DATA HDF5 FILE PATH>" \
    --train_data.transform_type="linear_prob" \
    --val_data.transform_type="test" \
    --load_checkpoint='' \
    --load_pretrained="<YOUR PRE-TRAINED MODEL PATH>"
```

Finetuning on ImageNet for both pre-trained MAE and M3AE
```
python3 -m m3ae.finetune_main \
    --seed=42 \
    --mae.model_type=large \
    --mae.drop_path=0.1 \
    --weight_decay=0.05 \
    --mixup_alpha=0.8 \
    --cutmix_alpha=1.0 \
    --switch_prob=0.5 \
    --label_smoothing=0.1 \
    --layer_decay=0.60 \
    --clip_gradient=1e9 \
    --batch_size=1024 \
    --warmup_epochs=5 \
    --epochs=100 \
    --dataloader_n_workers=16 \
    --dataloader_shuffle=False \
    --log_freq=500 \
    --save_model_freq=10000 \
    --lr_peak_value=1e-3 \
    --train_data.partition="train" \
    --val_data.partition="val" \
    --train_data.path="<YOUR DATA HDF5 FILE PATH>" \
    --val_data.path="<YOUR DATA HDF5 FILE PATH>" \
    --train_data.transform_type="finetune" \
    --val_data.transform_type="test" \
    --load_pretrained="<YOUR PRE-TRAINED MODEL PATH>"
```

## HDF5 Data Format
In order to facilitate training on cloud, we store all the dataset
as HDF5 files and read them from cloud storage buckets. For paired image and text
dataset, the HDF5 data contains two field, `jpg` and `caption`. The `jpg` field
is an 1D array containing the raw bytes of JPEG encoded images. The `caption`
field is an 1D array of utf-8 encoded text. For ImageNet dataset, the image JPEG
bytes are stored in field `train_jpg` and `val_jpg`, and the integer labels are
stored in field `train_labels` and `val_labels`. For unpaired text only dataset,
the utf-8 encoded text is stored in field `text`.

## Pre-trained Model Weights
[Pre-trained model weights can be downloaded here](https://www.dropbox.com/scl/fo/hnw0qtdfntl9mxls2xb8g/h?rlkey=vyzg1fm71prznifv6ijjy4t6m&dl=0).
The M3AE and MAE models here are trained for 50 epochs on the CC12M dataset using
the hyperparameters specified in the paper.

For converting the pre-trained Jax weights to PyTorch, please refer to [this colab](https://colab.research.google.com/drive/1L6EiUUSkeTzWtB3_K4fR0evn-SDqKWME).

## Credits
* The MAE is heavily inspired by the [original MAE implementation](https://github.com/facebookresearch/mae).

* The VQGAN image tokenizers are from [MaskGiT](https://github.com/google-research/maskgit)
and [dalle-mini](https://huggingface.co/dalle-mini/vqgan_imagenet_f16_16384).

* The transformer implementation is heavily inspired by [jax-models](https://github.com/DarshanDeshpande/jax-models).

* Some utilities are borrowed from [JaxCQL](https://github.com/young-geng/JaxCQL).

## Contact
If you have any questions, please open an issue or contact young.geng@berkeley.edu and hao.liu@berkely.edu.