細胞E_2

CELL-E 2：使用雙向文字到圖像轉換器將蛋白質轉換為圖片並返回

該存儲庫是 CELL-E 2：使用雙向文字到圖像轉換器將蛋白質轉換為圖片並返回的官方實現。

要求

建立虛擬環境並透過以下方式安裝所需的軟體包：

 pip install -r requirements.txt

接下來，使用適當的 CUDA 版本安裝torch = 2.0.0

模型下載

模型可在 Hugging Face 上找到。

我們還有兩個可用空間，您可以在其中對自己的數據進行預測！

• 影像預測
• 序列預測

產生影像

若要產生圖像，請將已儲存的模型設定為 ckpt_path。此方法可能不穩定，因此請參考Demo.ipynb查看另一種載入方式。

 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 )

序列預測

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

訓練

CELL-E 的訓練分為 3 個階段：

• 訓練蛋白質閾值圖像編碼器
• 訓練 Nucleus 影像編碼器
• 培訓 CELL-E 變壓器

VQGAN

如果使用蛋白質閾值影像，請為資料集設定threshold: True 。

我們使用馴服變壓器程式碼的稍微修改版本。

要訓練，請執行以下腳本：

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

請參閱原始儲存庫以取得其他標誌，例如--gpus 。

準備資料集

圖片

我們在腳本資料夾中提供了下載人類蛋白質圖譜和 OpenCell 影像的腳本。資料載入器需要data_csv 。您必須產生一個 csv 文件，其中包含nucleus_image_path 、 protein_image_path 、 metadata_path 、 split （train 或 val）和sequence （可選）列。假設該檔案與影像和元資料檔案存在於同一通用data資料夾中。

元數據

元資料是一個 JSON，應該伴隨每個蛋白質序列。如果序列未出現在data_csv中，則它必須出現在metadata.json中，並帶有名為protein_sequence的鍵。

此處添加更多資訊對於查詢單一蛋白​​質非常有用。它們可以透過retrieve_metadata檢索，這會在資料集物件中建立一個self.metadata變數。

要訓練，請執行以下腳本：

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

以與 VQGAN 相同的格式指定--gpus 。

CELL-E 包含以下選項：

• ckpt_path ：恢復先前的 CELL-E 2 訓練。使用 state_dict 儲存模型
• vqgan_model_path ：為蛋白質影像編碼器保存的蛋白質影像模型（帶有 state_dict）
• vqgan_config_path ：保存的蛋白質影像模型yaml
• condition_model_path ：為蛋白質影像編碼器保存的條件（核）模型（帶有 state_dict）
• condition_config_path ：保存的條件（核心）模型 yaml
• num_images ：如果僅使用蛋白質影像編碼器，則為 1；如果包含條件影像編碼器，則為 2
• image_key : nucleus 、 target或threshold
• dim : 語言模型嵌入的維度
• num_text_tokens ：語言模型中的標記總數（ESM-2 為 33）
• text_seq_len ：考慮的氨基酸總數
• depth ：Transformer 模型深度，通常越深越好，但會犧牲 VRAM
• heads ：多頭注意力中使用的頭數
• dim_head ：注意力頭的大小
• attn_dropout ：訓練中的注意力損失率
• ff_dropout ：訓練中的前饋 Dropout 率
• loss_img_weight ：應用於影像重建的加權。文字權重 = 1
• loss_text_weight ：應用於條件式影像重建的加權。
• stable ：規範權重（用於發生梯度爆炸時）
• learning_rate ：Adam 優化器的學習率
• monitor ：用於儲存模型的參數

引文

如果您決定將我們的程式碼用於您研究的任何部分，請引用我們。

 @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}
}