lhao499 / language-quantized-autoencoders

Language Quantized AutoEncoders

Language Quantized AutoEncoders

This is a Jax implementation of our work Language Quantized AutoEncoders.

It contains training and evalutation code.

This implementation has been tested on multi-GPU and Google Cloud TPU and supports both multi-host training with TPU Pods and multi-GPU training.

Usage

Experiments can be launched via the following commands.

An example script of launching a LQAE training job is:

export SCRIPT_DIR="$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )"
export PROJECT_DIR="$( cd -- "$( dirname -- "$SCRIPT_DIR" )" &> /dev/null && pwd )"
cd $PROJECT_DIR
export PYTHONPATH="$PYTHONPATH:$PROJECT_DIR"

echo $PYTHONPATH
export WANDB_API_KEY=''

export experiment_name='lqae'
export project_id='lqae'
export wu='5'
export ep='100'
export model='lqae'
export experiment_note=""
export experiment_id="lqae-base"

python3 -m lqae.main.lqae_main \
    --model_type="$model" \
    --lqae.bert_min_ratio=0.5 \
    --lqae.bert_max_ratio=0.5 \
    --lqae.quantizer_loss_commitment=0.005 \
    --lqae.quantizer_loss_entropy=0.0 \
    --lqae.quantizer_loss_perplexity=0.0 \
    --lqae.l2_normalize=True \
    --lqae.top_k_value=1 \
    --lqae.top_k_avg=False \
    --lqae.top_k_rnd=False \
    --lqae.vit_encoder_decoder=True \
    --lqae.vit_model_type='base' \
    --lqae.patch_size=16 \
    --lqae.use_bert_codebook=True \
    --lqae.bert_mask_loss_weight=0.0001 \
    --lqae.bert_channel_image_loss_weight=1.0 \
    --lqae.nochannel_image_loss_weight=0.0 \
    --lqae.quantizer_latent_dim=0 \
    --lqae.strawman_codebook=False \
    --lqae.use_bert_ste=False \
    --seed=42 \
    --epochs="$ep" \
    --lr_warmup_epochs="$wu" \
    --batch_size=512 \
    --dataloader_n_workers=16 \
    --log_freq=500 \
    --plot_freq=2000 \
    --save_model_freq=10000 \
    --lr_peak_value=1.5e-4 \
    --weight_decay=0.0005 \
    --load_checkpoint='' \
    --dataset='imagenet' \
    --imagenet_data.path="YOUR IMAGENET FILE in HDF5" \
    --imagenet_data.random_start=True \
    --log_all_worker=False \
    --logging.online=True \
    --logging.project_id="$project_id" \
    --logging.experiment_id="$experiment_id" \
    --logging.experiment_note="$experiment_note" \
    --logging.output_dir="$HOME/experiment_output/$project_id"

Example of running LLM based evaluation using LQAE pretrained model is at this colab.

To run experiments more conveniently on TPUs, you may want to use the script in jobs folder to manage TPUs jobs.

language-quantized-autoencoders's Issues

Environment setting

Hi, thank you for the interesting work!

I'm trying to reproduce your results using multi-gpu setting, however setting up the environment seems to have conflicts.

I tried to create the environment with

 conda env create -f gpu_requirement.yml 

Has anybody succeeded in setting up the gpu environment?

Pytorch version

Hi @lhao499 ,
Thank you for sharing this code.
Would you give the pytorch version for this code?

Thanks.

LQAE Model Checkpoint

Hi. Thank you for conducting interesting research.
Are you planning to release the checkpoint for the LQAE model?