sandeepusa/vermeer-test2

Vermeer

Autoregressive generative modeling of microscopy predicts protein localization.

Vermeer extends autoregressive generative modeling of natural images to multi-channel microscopy images. We train Vermeer on the Human Protein Atlas, using protein language model features of protein sequence as conditioning information.

Given the first three reference channels and a protein's ESM-C embedding, the model generates the corresponding fluorescent protein channel. Tutorial.ipynb demonstrates this virtual staining task. To run this notebook, first download the example data and example model checkpoints (contains tokenizer checkpoint and vermeer-L checkpoint) from vermeer_models_and_examples.zip. The provided yaml file ar_microscopy_gen_03032026.yaml can be used to create a conda environment to run this notebook.

Training a new model

To train a new model, first download the raw microscopy images from the Human Protein Atlas using the scripts code/download_preprocess_data/download_images_parallel.py and code/download_preprocess_data/hpa_stratified_preprocessing_final.py:

python download_images_parallel.py --output-dir <OUTPUT_DIR>
python hpa_stratified_preprocessing_final.py --image_size 256
# specify the paths in the config file at the beginning of this script

Then tokenize all of the images and compute the ESM-C embeddings:

python dataset/prepare_protein_prefix.py \
    --input_dir input_dir \
    --h5_filename protein_prefix.h5 \
    --metadata_dir output_metadir \
    --device cuda

python autoregressive/train/extract_codes_ca.py \
    --data-path input_data_path \
    --code-path output_code_path \
    --vq-ckpt pretrained_models/vq_ds16_c2i.pt \
    --ten-crop \
    --rotate \
    --debug \
    --label-type esm_embed_mean_pool \
    --label-file protein_prefix.h5 \
    --n-channels 4 \
    --num-workers 4 \
    --image-size 256

Download the pretrained LlamaGen checkpoint from LlamaGen (e.g. c2i_L_256.pt). Then run the training script code/CA_LlamaGen/autoregressive/train/train_ca.py:

torchrun autoregressive/train/train_ca.py\
    --nnodes=1 --nproc_per_node=2 --node_rank=0 \
    --master_port=12334 \
    autoregressive/train/train_ca.py \
    --results-dir output_dir \
    --val-dirs val1,val2,cell_line_holdouts \
    --code-path output_code_path \
    --image-size 256 \
    --gpt-model GPT-L \
    --num-workers 8 \
    --ckpt-every 5000 \
    --lr 1e-4 \
    --epochs 150 \
    --experiment-name "hpa_split_size_256_lr_1e-4_lin_schedule" \
    --global-batch-size 96 \
    --val-every 1000 \
    --gpt-type ca_esm_embed_mean_pool \
    --pretrained-gpt-ckpt $pretrained_gpt_ckpt_dir/c2i_L_256.pt \
    --lr-schedule lin \
    --warmup-epochs 10

Evaluation can be run using the evaluation script, code/CA_LlamaGen/scripts/eval/run_eval_pipeline.sh.

Acknowledgements

Built on LlamaGen. Protein representations from ESM-C. Microscopy data from the Human Protein Atlas.