细胞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}
}