CELDA E_2

CELL-E 2: Traducir proteínas a imágenes y viceversa con un transformador bidireccional de texto a imagen

Este repositorio es la implementación oficial de CELL-E 2: Traducción de proteínas a imágenes y viceversa con un transformador bidireccional de texto a imagen.

Requisitos

Cree un entorno virtual e instale los paquetes necesarios a través de:

 pip install -r requirements.txt

A continuación, instale torch = 2.0.0 con la versión CUDA adecuada.

Descargas de modelos

Los modelos están disponibles en Hugging Face.

¡También tenemos dos espacios disponibles donde puedes ejecutar predicciones sobre tus propios datos!

• Predicción de imágenes
• Predicción de secuencia

Generando imágenes

Para generar imágenes, configure el modelo guardado como ckpt_path. Este método puede ser inestable, así que consulte Demo.ipynb para ver otra forma de carga.

 from omegaconf import OmegaConf
from celle_main import instantiate_from_config

configs = OmegaConf . load ( configs / celle . yaml );

model = instantiate_from_config ( configs . model ). to ( device );

model . sample ( text = sequence ,
                      condition = nucleus ,
                      return_logits = True ,
                      progress = True )

Predicción de secuencia

 model . sample_text ( condition = nucleus ,
                  image = image ,
                  return_logits = True ,
                  progress = True )

Capacitación

El entrenamiento para CELL-E se produce en 3 etapas:

• Codificador de imagen de umbral de proteína de entrenamiento
• Entrenamiento de un codificador de imágenes Nucleus
• Transformador CELL-E de formación

VQGAN

Si utiliza la imagen del umbral de proteínas, establezca threshold: True para el conjunto de datos.

Usamos una versión ligeramente modificada del código de domesticación de transformadores.

Para entrenar, ejecute el siguiente script:

python celle_taming_main.py --base configs/threshold_vqgan.yaml -t True

Consulte el repositorio original para ver indicadores adicionales, como --gpus .

Preparando el conjunto de datos

Imágenes

Proporcionamos secuencias de comandos para descargar imágenes de Human Protein Atlas y OpenCell en la carpeta de secuencias de comandos. Se necesita un data_csv para el cargador de datos. Debe generar un archivo csv que contenga las columnas nucleus_image_path , protein_image_path , metadata_path , split (train o val) y sequence (opcional). Se supone que este archivo existe dentro de la misma carpeta data generales que los archivos de imágenes y metadatos.

Metadatos

Los metadatos son un JSON que deben acompañar a cada secuencia de proteínas. Si una secuencia no aparece en data_csv , debe aparecer en metadata.json con una clave denominada protein_sequence .

Agregar más información aquí puede resultar útil para consultar proteínas individuales. Se pueden recuperar a través de retrieve_metadata , que crea una variable self.metadata dentro del objeto del conjunto de datos.

Para entrenar, ejecute el siguiente script:

python celle_main.py --base configs/celle.yaml -t True

Especifique --gpus en el mismo formato que VQGAN.

CELL-E contiene las siguientes opciones:

• ckpt_path : reanuda el entrenamiento anterior de CELL-E 2. Modelo guardado con state_dict
• vqgan_model_path : modelo de imagen de proteína guardado (con state_dict) para el codificador de imágenes de proteína
• vqgan_config_path : modelo de imagen de proteína guardada yaml
• condition_model_path : modelo de condición (núcleo) guardado (con state_dict) para el codificador de imágenes de proteínas
• condition_config_path : modelo de condición guardada (núcleo) yaml
• num_images : 1 si solo se usa el codificador de imágenes de proteínas, 2 si se incluye el codificador de imágenes de condición
• image_key : nucleus , target o threshold
• dim : dimensión de la incrustación del modelo de lenguaje
• num_text_tokens : número total de tokens en el modelo de lenguaje (33 para ESM-2)
• text_seq_len : Número total de aminoácidos considerados
• depth : profundidad del modelo de transformador, más profundo suele ser mejor a costa de VRAM
• heads : número de cabezas utilizadas en atención multicabezal
• dim_head : tamaño de las cabezas de atención
• attn_dropout : Tasa de abandono de la atención en el entrenamiento
• ff_dropout : Tasa de abandono de feed-forward en el entrenamiento
• loss_img_weight : Ponderación aplicada a la reconstrucción de imágenes. peso del texto = 1
• loss_text_weight : ponderación aplicada a la reconstrucción de imágenes condicionadas.
• stable : Pesos de normas (para cuando se producen gradientes explosivos)
• learning_rate : tasa de aprendizaje para el optimizador Adam
• monitor : parámetro utilizado para guardar modelos

Citación

Cítenos si decide utilizar nuestro código para cualquier parte de su investigación.

 @inproceedings{
anonymous2023translating,
title={CELL-E 2: Translating Proteins to Pictures and Back with a Bidirectional Text-to-Image Transformer},
author={Emaad Khwaja, Yun S. Song, Aaron Agarunov, and Bo Huang},
booktitle={Thirty-seventh Conference on Neural Information Processing Systems},
year={2023},
url={https://openreview.net/forum?id=YSMLVffl5u}
}